[Reverted][RL] Support Rollout Routing Replay

fastdeploy/config.py

@@ -1484,6 +1484,31 @@ def __str__(self) -> str:
         return json.dumps({key: value for key, value in self.__dict__.items()})
 
 
+class RoutingReplayConfig:
+    """Configuration for Routing Replay used in RL training"""
+
+    def __init__(self, args) -> None:
+        self.enable_routing_replay: bool = False
+        self.routing_store_type: str = "local"
+
+        # Local routing store
+        self.local_store_dir: str = "./routing_replay_output"
+
+        # RDMA routing store
+        # TODO: Add RDMA routing store configuration attributes here when the feature is implemented.
+
+        if args is not None:
+            for key, value in args.items():
+                if hasattr(self, key) and value != "None":
+                    setattr(self, key, value)
+
+    def to_json_string(self):
+        """
+        Convert routing replay config to json string.
+        """
+        return json.dumps({key: value for key, value in self.__dict__.items()})
+
+
 class FDConfig:
     """
     The configuration class which contains all fastdeploy-related configuration. This
@@ -1517,6 +1542,7 @@ def __init__(
         early_stop_config: Optional[Dict[str, Any]] = None,
         tool_parser: str = None,
         test_mode=False,
+        routing_replay_config: Optional[RoutingReplayConfig] = None,
     ):
         self.model_config: ModelConfig = model_config  # type: ignore
         self.cache_config: CacheConfig = cache_config  # type: ignore
@@ -1533,6 +1559,7 @@ def __init__(
         self.plas_attention_config: Optional[PlasAttentionConfig] = plas_attention_config
         self.structured_outputs_config: StructuredOutputsConfig = structured_outputs_config
         self.router_config: RouterConfig = router_config
+        self.routing_replay_config = routing_replay_config
 
         # Initialize cuda graph capture list
         max_capture_shape = self.scheduler_config.max_num_seqs

fastdeploy/engine/args_utils.py

@@ -35,6 +35,7 @@
     PlasAttentionConfig,
     PoolerConfig,
     RouterConfig,
+    RoutingReplayConfig,
     RunnerOption,
     SpeculativeConfig,
     StructuredOutputsConfig,
@@ -491,6 +492,11 @@ class EngineArgs:
     Configuration for eplb.
     """
 
+    routing_replay_config: Optional[Dict[str, Any]] = None
+    """
+    Flag to rollout routing replay(r3)
+    """
+
     def __post_init__(self):
         """
         Post-initialization processing to set default tokenizer if not provided.
@@ -884,6 +890,12 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             default=EngineArgs.eplb_config,
             help="Config of eplb.",
         )
+        parallel_group.add_argument(
+            "--routing-replay-config",
+            type=json.loads,
+            default=EngineArgs.routing_replay_config,
+            help="Flag of rollout routing replay(r3).",
+        )
         parallel_group.add_argument(
             "--enable-chunked-moe",
             action="store_true",
@@ -1237,6 +1249,14 @@ def create_eplb_config(self) -> EPLBConfig:
         eplb_args["enable_eplb"] = self.enable_eplb
         return EPLBConfig(eplb_args)
 
+    def create_routing_repaly_config(self) -> RoutingReplayConfig:
+        """ """
+        routing_replay_args = asdict(self)
+        if self.routing_replay_config is not None:
+            for k, v in self.routing_replay_config.items():
+                routing_replay_args[k] = v
+        return RoutingReplayConfig(routing_replay_args)
+
     def create_engine_config(self, port_availability_check=True) -> FDConfig:
         """
         Create and return a Config object based on the current settings.
@@ -1280,6 +1300,7 @@ def create_engine_config(self, port_availability_check=True) -> FDConfig:
         graph_opt_cfg = self.create_graph_optimization_config()
         plas_attention_config = self.create_plas_attention_config()
         eplb_cfg = self.create_eplb_config()
+        routing_replay_config = self.create_routing_repaly_config()
         router_config = RouterConfig(all_dict)
 
         early_stop_cfg = self.create_early_stop_config()
@@ -1312,4 +1333,5 @@ def create_engine_config(self, port_availability_check=True) -> FDConfig:
             graph_opt_config=graph_opt_cfg,
             plas_attention_config=plas_attention_config,
             early_stop_config=early_stop_cfg,
+            routing_replay_config=routing_replay_config,
         )

fastdeploy/engine/engine.py

@@ -568,6 +568,7 @@ def _start_worker_service(self):
             f" --logprobs_mode {self.cfg.model_config.logprobs_mode}"
             f" --max_logprobs {self.cfg.model_config.max_logprobs}"
             f" --eplb_config '{self.cfg.eplb_config.to_json_string()}'"
+            f" --routing_replay_config '{self.cfg.routing_replay_config.to_json_string()}'"
         )
         if self.cfg.structured_outputs_config.logits_processors is not None:
             arguments += f" --logits-processors {' '.join(self.cfg.structured_outputs_config.logits_processors)}"

fastdeploy/model_executor/forward_meta.py

@@ -142,6 +142,8 @@ class ForwardMeta:
     caches: Optional[list[paddle.Tensor]] = None
     # Flag of profile run
     is_dummy_or_profile_run: bool = False
+    # Routing Replay table buffer
+    routing_replay_table: Optional[paddle.Tensor] = None
 
     # chunked MoE related
     moe_num_chunk: int = 1

fastdeploy/model_executor/layers/backends/dcu/fused_moe_triton_backends.py

@@ -14,6 +14,8 @@
 # limitations under the License.
 """
 
+from typing import Callable
+
 import paddle
 from paddle import nn
 
@@ -101,6 +103,7 @@ def apply(
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
+        topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
         Triton compute Fused MoE.
@@ -117,6 +120,8 @@ def apply(
         scores += layer.gate_correction_bias
         topk_weights, topk_ids = paddle.topk(scores, k=top_k, axis=-1, sorted=False)
         topk_weights = topk_weights / topk_weights.sum(axis=-1, keepdim=True)
+        if topk_ids_hookfunc is not None:
+            topk_ids_hookfunc(topk_ids=topk_ids)
 
         intermediate_cache1 = paddle.empty(
             [token_num * top_k, moe_intermediate_size * 2],

fastdeploy/model_executor/layers/backends/gcu/moe/fused_moe_method_gcu_backend.py

@@ -16,6 +16,7 @@
 
 import multiprocessing
 import os
+from typing import Callable
 
 import numpy as np
 import paddle
@@ -182,6 +183,7 @@ def apply(
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
+        topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
         Paddle gcu compute Fused MoE.
@@ -194,6 +196,7 @@ def apply_ep_prefill(
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
+        topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
         Apply the EP prefill method.
@@ -205,6 +208,7 @@ def apply_ep_decode(
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
+        topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
         Apply the EP decoder method.
@@ -216,6 +220,7 @@ def apply_tp(
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
+        topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
         Paddle Cutlass compute Fused MoE.
@@ -381,6 +386,7 @@ def apply(
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
+        topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
         Paddle gcu compute Fused MoE.

fastdeploy/model_executor/layers/backends/metax/moe/fused_moe_triton_metax_backend.py

@@ -14,6 +14,8 @@
 # limitations under the License.
 """
 
+from typing import Callable
+
 import paddle
 from paddle import nn
 
@@ -245,6 +247,7 @@ def apply(
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
+        topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
         Triton compute Fused MoE.
@@ -274,6 +277,9 @@ def apply(
                 True,  # apply_norm_weight
                 False,
             )
+        if topk_ids_hookfunc is not None:
+            topk_ids_hookfunc(topk_ids=topk_ids)
+
         up_gate_proj_out = paddle.empty(
             [token_num * top_k, moe_intermediate_size * 2],
             dtype=x.dtype,

fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py

@@ -15,6 +15,7 @@
 """
 
 from abc import abstractmethod
+from typing import Callable
 
 import paddle
 from paddle import nn
@@ -163,6 +164,7 @@ def apply_ep_prefill(
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
+        topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
         Apply the EP prefill method.
@@ -175,6 +177,7 @@ def apply_ep_decode(
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
+        topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
         Apply the EP decoder method.
@@ -187,6 +190,7 @@ def apply_tp(
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
+        topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
         Paddle Cutlass compute Fused MoE.
@@ -198,6 +202,7 @@ def apply(
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
+        topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
         Paddle Cutlass compute Fused MoE.
@@ -207,13 +212,13 @@ def apply(
             if layer.fd_config.model_config.moe_phase.phase == "prefill":
                 if layer.fd_config.scheduler_config.splitwise_role == "mixed" and is_moe_start_layer:
                     self.ep_prefill_runner.clean_low_latency_buffer()
-                return self.apply_ep_prefill(layer, x, gate)
+                return self.apply_ep_prefill(layer, x, gate, topk_ids_hookfunc=topk_ids_hookfunc)
             else:
                 if layer.fd_config.scheduler_config.splitwise_role == "mixed" and is_moe_start_layer:
                     self.ep_decoder_runner.clean_low_latency_buffer()
-                return self.apply_ep_decode(layer, x, gate)
+                return self.apply_ep_decode(layer, x, gate, topk_ids_hookfunc=topk_ids_hookfunc)
         else:
-            return self.apply_tp(layer, x, gate)
+            return self.apply_tp(layer, x, gate, topk_ids_hookfunc=topk_ids_hookfunc)
 
 
 class UnquantizedFusedMoEMethod(MoEMethodBase):

fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py

@@ -14,6 +14,8 @@
 # limitations under the License.
 """
 
+from typing import Callable
+
 import paddle
 from paddle import nn
 from paddle.nn.quant import weight_quantize
@@ -132,6 +134,7 @@ def apply_ep_prefill(
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
+        topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
         Apply the EP prefill method.
@@ -148,8 +151,13 @@ def apply_ep_prefill(
             handle,
             event,
         ) = self.ep_prefill_runner.dispatch(x, topk_idx, topk_weights)
+
+        if topk_ids_hookfunc is not None:
+            topk_ids_hookfunc(topk_ids=topk_idx)
+
         if self.ep_prefill_runner.ep_engine.async_finish:
             event.current_stream_wait()
+
         token_all_num = sum(recv_num_tokens_per_expert_list)
 
         # 3. Compute ffn
@@ -217,6 +225,7 @@ def apply_ep_decode(
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
+        topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
         Apply the EP decoder method.
@@ -225,6 +234,10 @@ def apply_ep_decode(
         estimate_total_token_nums = gate_out.shape[0] * layer.top_k
         # 1. Select topk experts and weights
         topk_idx, topk_weights = self.ep_decoder_runner.moe_select(layer, gate_out)
+
+        if topk_ids_hookfunc is not None:
+            topk_ids_hookfunc(topk_ids=topk_idx)
+
         expertwise_scale = None
         if hasattr(layer, "up_gate_proj_in_scale_all_experts"):  # only use in w4a8
             expertwise_scale = getattr(layer, "up_gate_proj_in_scale_all_experts", None)
@@ -269,6 +282,7 @@ def apply_tp(
         layer: nn.Layer,
         x: paddle.Tensor,
         gate: nn.Layer,
+        topk_ids_hookfunc: Callable = None,
     ) -> paddle.Tensor:
         """
         Paddle Cutlass compute Fused MoE.
@@ -369,6 +383,9 @@ def apply_tp(
         if hasattr(layer, "up_gate_proj_in_scale"):
             dequant_scale = None
 
+        if topk_ids_hookfunc is not None:
+            topk_ids_hookfunc(topk_ids=topk_idx)
+
         if not layer.with_bias and self.moe_quant_type != "w4a8" and self.moe_quant_type != "w4afp8":
             # only w4a8 need expert_idx_per_token
             # Other need not this tensor, so we make it None.