Merge remote-tracking branch 'origin/develop' into develop

Author: TBD1

custom_ops/utils/auto_gen_w4afp8_gemm_kernel.py

@@ -94,6 +94,8 @@
     [2560, 1536, 64, 0, 128],
     [1536, 2560, 64, 0, 128],
     [2560, 768, 64, 0, 128],
+    [768, 2048, 128, 0, 128],
+    [2048, 384, 128, 0, 128],
 ]
 
 dtype = ["BF16"]

examples/mooncake_store/README.md examples/cache_storage/README.mdexamples/mooncake_store/README.md renamed to examples/cache_storage/README.md

@@ -4,6 +4,6 @@
     "global_segment_size":8589934592,
     "local_buffer_size":134217728,
     "protocol":"rdma",
-    "rdma_devices": "mlx5_1,mlx5_2,mlx5_3,mlx5_4",
+    "rdma_devices": "",
     "master_server_addr":"0.0.0.0:15001"
 }

…ples/mooncake_store/mooncake_config.json …mples/cache_storage/mooncake_config.jsonexamples/mooncake_store/mooncake_config.json renamed to examples/cache_storage/mooncake_config.json examples/mooncake_store/mooncake_config.json renamed to examples/cache_storage/mooncake_config.json

@@ -17,7 +17,6 @@ S1_PORT=52800
 ports=(
     $S0_PORT $((S0_PORT + 1)) $((S0_PORT + 2)) $((S0_PORT + 3))
     $S1_PORT $((S1_PORT + 1)) $((S1_PORT + 2)) $((S1_PORT + 3))
-    $ROUTER_PORT
 )
 check_ports "${ports[@]}" || {
     echo "❌ Some ports are in use. Please release them."

examples/mooncake_store/run.sh examples/cache_storage/run.shexamples/mooncake_store/run.sh renamed to examples/cache_storage/run.sh

@@ -265,40 +265,41 @@ def launch_cache_manager(
         else:
             kvcache_storage_backend_str = "none"
 
-        for i in range(tensor_parallel_size):
-            launch_cmd = (
-                "FLAGS_allocator_strategy=auto_growth "
-                + visible_devices
-                + " NCCL_MAX_NCHANNELS=1 NCCL_BUFFSIZE=0"
-                + f" FD_ENABLE_SWAP_SPACE_CLEARING={envs.FD_ENABLE_SWAP_SPACE_CLEARING}"
-                + f" {sys.executable} {py_path}"
-                + f" --device_id {int(device_ids[i])}"
-                + f" --rank {i}"
-                + f" --splitwise_role {self.splitwise_role}"
-                + f" --num_layers {cache_config.model_cfg.num_hidden_layers}"
-                + f" --mp_num {tensor_parallel_size}"
-                + f" --cache_dtype {cache_config.cache_dtype}"
-                + f" --key_cache_shape {key_cache_shape}"
-                + val_cache_arg_str
-                + f" --cache_queue_port {cache_config.local_cache_queue_port}"
-                + f" --enable_splitwise {int(self.enable_splitwise)}"
-                + f" --pod_ip {pod_ip}"
-                + f" --engine_worker_queue_port {engine_worker_queue_port}"
-                + f" --num_cpu_blocks {cache_config.num_cpu_blocks}"
-                + f" --ipc_suffix {ipc_suffix}"
-                + f" --protocol {cache_config.cache_transfer_protocol}"
-                + f" --local_data_parallel_id {self.local_data_parallel_id}"
-                + f" --rdma_port {cache_config.local_rdma_comm_ports[i] if cache_config.local_rdma_comm_ports is not None else '0'}"
-                + f" --speculative_config '{self.speculative_config.to_json_string()}'"
-                + f" --default_dtype '{self.config.model_config.dtype}'"
-                + (" --create_cache_tensor" if create_cache_tensor else "")
-                + f" --kvcache_storage_backend {kvcache_storage_backend_str}"
-                + f" --write_policy {cache_config.write_policy}"
-                + f" --max_model_len {self.config.model_config.max_model_len}"
-                + f" >{log_dir}/launch_cache_transfer_manager_{int(device_ids[i])}.log 2>&1"
-            )
-            logger.info(f"Launch cache transfer manager, command:{launch_cmd}")
-            cache_manager_processes.append(subprocess.Popen(launch_cmd, shell=True, preexec_fn=os.setsid))
+        if self.cache_config.swap_space or self.cache_config.kvcache_storage_backend:
+            for i in range(tensor_parallel_size):
+                launch_cmd = (
+                    "FLAGS_allocator_strategy=auto_growth "
+                    + visible_devices
+                    + " NCCL_MAX_NCHANNELS=1 NCCL_BUFFSIZE=0"
+                    + f" FD_ENABLE_SWAP_SPACE_CLEARING={envs.FD_ENABLE_SWAP_SPACE_CLEARING}"
+                    + f" {sys.executable} {py_path}"
+                    + f" --device_id {int(device_ids[i])}"
+                    + f" --rank {i}"
+                    + f" --splitwise_role {self.splitwise_role}"
+                    + f" --num_layers {cache_config.model_cfg.num_hidden_layers}"
+                    + f" --mp_num {tensor_parallel_size}"
+                    + f" --cache_dtype {cache_config.cache_dtype}"
+                    + f" --key_cache_shape {key_cache_shape}"
+                    + val_cache_arg_str
+                    + f" --cache_queue_port {cache_config.local_cache_queue_port}"
+                    + f" --enable_splitwise {int(self.enable_splitwise)}"
+                    + f" --pod_ip {pod_ip}"
+                    + f" --engine_worker_queue_port {engine_worker_queue_port}"
+                    + f" --num_cpu_blocks {cache_config.num_cpu_blocks}"
+                    + f" --ipc_suffix {ipc_suffix}"
+                    + f" --protocol {cache_config.cache_transfer_protocol}"
+                    + f" --local_data_parallel_id {self.local_data_parallel_id}"
+                    + f" --rdma_port {cache_config.local_rdma_comm_ports[i] if cache_config.local_rdma_comm_ports is not None else '0'}"
+                    + f" --speculative_config '{self.speculative_config.to_json_string()}'"
+                    + f" --default_dtype '{self.config.model_config.dtype}'"
+                    + (" --create_cache_tensor" if create_cache_tensor else "")
+                    + f" --kvcache_storage_backend {kvcache_storage_backend_str}"
+                    + f" --write_policy {cache_config.write_policy}"
+                    + f" --max_model_len {self.config.model_config.max_model_len}"
+                    + f" >{log_dir}/launch_cache_transfer_manager_{int(device_ids[i])}.log 2>&1"
+                )
+                logger.info(f"Launch cache transfer manager, command:{launch_cmd}")
+                cache_manager_processes.append(subprocess.Popen(launch_cmd, shell=True, preexec_fn=os.setsid))
 
         logger.info("PrefixCacheManager is waiting for kv cache to be initialized.")
         while np.sum(self.cache_ready_signal.value) != tensor_parallel_size:
@@ -308,13 +309,14 @@ def launch_cache_manager(
             while np.sum(self.swap_space_ready_signal.value) != tensor_parallel_size:
                 time.sleep(1)
 
-        exit_code = cache_manager_processes[-1].poll()
-        if exit_code is None:
-            logger.info("Launch cache transfer manager successful")
-        else:
-            logger.info(
-                "Launch cache transfer manager failed, see launch_cache_transfer_manager.log for more information"
-            )
+        if cache_manager_processes:
+            exit_code = cache_manager_processes[-1].poll()
+            if exit_code is None:
+                logger.info("Launch cache transfer manager successful")
+            else:
+                logger.info(
+                    "Launch cache transfer manager failed, see launch_cache_transfer_manager.log for more information"
+                )
 
         # Start additional threads
         if cache_config.kvcache_storage_backend or self.num_cpu_blocks > 0:

examples/mooncake_store/utils.sh examples/cache_storage/utils.shexamples/mooncake_store/utils.sh renamed to examples/cache_storage/utils.sh

@@ -50,10 +50,10 @@ def create() -> "MooncakeStoreConfig":
         file_path = os.getenv("MOONCAKE_CONFIG_PATH")
 
         if file_path is None:
-            local_hostname = os.environ.get("MOONCAKE_LOCAL_HOSTNAME")
+            local_hostname = os.environ.get("MOONCAKE_LOCAL_HOSTNAME", "localhost")
             metadata_server = os.environ.get("MOONCAKE_METADATA_SERVER")
-            global_segment_size = os.environ.get("MOONCAKE_GLOBAL_SEGMENT_SIZE", DEFAULT_GLOBAL_SEGMENT_SIZE)
-            local_buffer_size = os.environ.get("MOONCAKE_LOCAL_BUFFER_SIZE", DEFAULT_LOCAL_BUFFER_SIZE)
+            global_segment_size = int(os.environ.get("MOONCAKE_GLOBAL_SEGMENT_SIZE", DEFAULT_GLOBAL_SEGMENT_SIZE))
+            local_buffer_size = int(os.environ.get("MOONCAKE_LOCAL_BUFFER_SIZE", DEFAULT_LOCAL_BUFFER_SIZE))
             protocol = os.environ.get("MOONCAKE_PROTOCOL", "rdma")
             rdma_devices = os.environ.get("MOONCAKE_RDMA_DEVICES", "")
             master_server_addr = os.environ.get("MOONCAKE_MASTER_SERVER_ADDR")
@@ -63,10 +63,10 @@ def create() -> "MooncakeStoreConfig":
             with open(file_path) as fin:
                 config = json.load(fin)
 
-            local_hostname = config.get("local_hostname")
+            local_hostname = config.get("local_hostname", "localhost")
             metadata_server = config.get("metadata_server")
-            global_segment_size = config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE)
-            local_buffer_size = config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)
+            global_segment_size = int(config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE))
+            local_buffer_size = int(config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE))
             protocol = config.get("protocol", "rdma")
             rdma_devices = config.get("rdma_devices", "")
             master_server_addr = config.get("master_server_addr")
@@ -75,6 +75,8 @@ def create() -> "MooncakeStoreConfig":
             # FIXME: use auto-select NICs in MooncakeStore will raise error and roll back to using TCP
             rdma_devices = get_rdma_nics()
             logger.info(f"No RDMA devices specified, defaulting to all available devices: {rdma_devices}")
+        if metadata_server is None or master_server_addr is None:
+            raise ValueError("Both MOONCAKE_METADATA_SERVER and MOONCAKE_MASTER_SERVER_ADDR must be provided.")
 
         return MooncakeStoreConfig(
             local_hostname=local_hostname,

fastdeploy/cache_manager/prefix_cache_manager.py

@@ -34,7 +34,6 @@
 from fastdeploy.input.preprocess import InputPreprocessor
 from fastdeploy.inter_communicator import (
     IPCSignal,
-    KVCacheStatus,
     ModelWeightsStatus,
     PrefixTreeStatus,
     RearrangeExpertStatus,
@@ -529,6 +528,19 @@ def update_model_weight(self, timeout=300):
         2 : worker update finish and notify client
         """
         with self.clear_update_lock:
+            if self.fd_config.cache_config.swap_space:
+                return False, "hierarchical cache updating is not supported"
+
+            if self.enable_prefix_caching:
+                # prefix_tree_status_signal: CLEARED -> UPDATING -> NORMAL
+                if self.prefix_tree_status_signal.value[0] == PrefixTreeStatus.CLEARED:
+                    self.prefix_tree_status_signal.value[0] = PrefixTreeStatus.UPDATING
+                    api_server_logger.info(f"Start to update prefix tree {self.prefix_tree_status_signal.value[0]}")
+                    while self.prefix_tree_status_signal.value[0] != PrefixTreeStatus.NORMAL:
+                        api_server_logger.info(f"..updating prefix tree {self.prefix_tree_status_signal.value[0]}")
+                        time.sleep(1)
+
+            # model_weights_status_signal: CLEARED -> UPDATING -> NORMAL
             if self.model_weights_status_signal.value[0] == ModelWeightsStatus.NORMAL:
                 return True, ""
             if self.model_weights_status_signal.value[0] == ModelWeightsStatus.UPDATING:
@@ -537,34 +549,13 @@ def update_model_weight(self, timeout=300):
                 return False, "worker is clearing model weight, cannot update now"
 
             self.model_weights_status_signal.value[0] = ModelWeightsStatus.UPDATING
-            if self.enable_prefix_caching or self.enable_splitwise:
-                self.kv_cache_status_signal.value[0] = KVCacheStatus.UPDATING
-            if self.enable_prefix_caching:
-                self.prefix_tree_status_signal.value[0] = PrefixTreeStatus.UPDATING
-            api_server_logger.info(f"start update model weight {self.model_weights_status_signal.value}")
-            all_updated = False
-            while timeout >= 0 and not all_updated:
-                api_server_logger.info(
-                    f"Updating model weights.. "
-                    f"model_weights_status: {self.model_weights_status_signal.value[0]}, "
-                    f"prefix_tree_status: {self.prefix_tree_status_signal.value[0]}, "
-                    f"kv_cache_status: {self.kv_cache_status_signal.value[0]} "
-                )
-                weight_updated = self.model_weights_status_signal.value[0] == ModelWeightsStatus.NORMAL
-                cache_updated = self.kv_cache_status_signal.value[0] == KVCacheStatus.NORMAL
-                prefix_updated = self.prefix_tree_status_signal.value[0] == PrefixTreeStatus.NORMAL
-                if self.enable_prefix_caching or self.enable_splitwise:
-                    if self.enable_prefix_caching:
-                        all_updated = weight_updated and cache_updated and prefix_updated
-                    else:
-                        all_updated = weight_updated and cache_updated
-                else:
-                    all_updated = weight_updated
+            api_server_logger.info(f"Start to update model weight {self.model_weights_status_signal.value[0]}")
+            while timeout >= 0 and self.model_weights_status_signal.value[0] != ModelWeightsStatus.NORMAL:
+                api_server_logger.info(f"..updating model weights {self.model_weights_status_signal.value[0]}")
                 time.sleep(1)
                 timeout -= 1
             if timeout < 0:
                 return False, "Update model weight timeout"
-            time.sleep(1)
             return True, ""
 
     def clear_load_weight(self, timeout=300):
@@ -575,6 +566,19 @@ def clear_load_weight(self, timeout=300):
         """
 
         with self.clear_update_lock:
+            if self.fd_config.cache_config.swap_space:
+                return False, "hierarchical cache clearing is not supported"
+
+            if self.enable_prefix_caching:
+                # prefix_tree_status_signal: NORMAL -> CLEARING -> CLEARED
+                if self.prefix_tree_status_signal.value[0] == PrefixTreeStatus.NORMAL:
+                    self.prefix_tree_status_signal.value[0] = PrefixTreeStatus.CLEARING
+                    api_server_logger.info(f"Start to clear prefix tree {self.prefix_tree_status_signal.value[0]}")
+                    while self.prefix_tree_status_signal.value[0] != PrefixTreeStatus.CLEARED:
+                        api_server_logger.info(f"..clearing prefix tree {self.prefix_tree_status_signal.value[0]}")
+                        time.sleep(1)
+
+            # model_weights_status_signal: NORMAL -> CLEARING -> CLEARED
             if self.model_weights_status_signal.value[0] == ModelWeightsStatus.CLEARED:
                 return True, ""
             if self.model_weights_status_signal.value[0] == ModelWeightsStatus.CLEARING:
@@ -583,36 +587,13 @@ def clear_load_weight(self, timeout=300):
                 return False, "worker is updating model weight, cannot clear now"
 
             self.model_weights_status_signal.value[0] = ModelWeightsStatus.CLEARING
-            if self.enable_prefix_caching or self.enable_splitwise:
-                self.kv_cache_status_signal.value[0] = KVCacheStatus.CLEARING
-            if self.enable_prefix_caching:
-                self.prefix_tree_status_signal.value[0] = PrefixTreeStatus.CLEARING
-
-            api_server_logger.info(f"start clear model weight {self.model_weights_status_signal.value}")
-            all_cleared = False
-            while timeout >= 0 and not all_cleared:
-                api_server_logger.info(
-                    f"Clearing model weights.. "
-                    f"model_weights_status: {self.model_weights_status_signal.value[0]}, "
-                    f"prefix_tree_status: {self.prefix_tree_status_signal.value[0]}, "
-                    f"kv_cache_status: {self.kv_cache_status_signal.value[0]} "
-                )
-                weight_cleared = self.model_weights_status_signal.value[0] == ModelWeightsStatus.CLEARED
-                cache_cleared = self.kv_cache_status_signal.value[0] == KVCacheStatus.CLEARED
-                prefix_cleared = self.prefix_tree_status_signal.value[0] == PrefixTreeStatus.CLEARED
-                if self.enable_prefix_caching or self.enable_splitwise:
-                    if self.enable_prefix_caching:
-                        all_cleared = weight_cleared and cache_cleared and prefix_cleared
-                    else:
-                        all_cleared = weight_cleared and cache_cleared
-                else:
-                    all_cleared = weight_cleared
+            api_server_logger.info(f"Start to clear model weight {self.model_weights_status_signal.value[0]}")
+            while timeout >= 0 and self.model_weights_status_signal.value[0] != ModelWeightsStatus.CLEARED:
+                api_server_logger.info(f"..clearing model weights {self.model_weights_status_signal.value[0]}")
                 time.sleep(1)
                 timeout -= 1
-
             if timeout < 0:
                 return False, "Clear model weight timeout"
-            time.sleep(1)
             return True, ""
 
     def check_model_weight_status(self):

fastdeploy/cache_manager/transfer_factory/mooncake_store/mooncake_store.py

@@ -546,7 +546,9 @@ async def completion_stream_generator(
                         reasoning_content="",
                         arrival_time=arrival_time,
                         logprobs=logprobs_res,
-                        prompt_logprobs=clamp_prompt_logprobs(prompt_logprobs_res),
+                        prompt_logprobs=(
+                            clamp_prompt_logprobs(prompt_logprobs_res) if not request.return_token_ids else None
+                        ),
                         draft_logprobs=draft_logprobs_res,
                         speculate_metrics=output_speculate_metrics,
                     )

fastdeploy/entrypoints/engine_client.py

@@ -36,6 +36,19 @@
 ]
 
 
+def _compute_hadamard_block_size(moe_intermediate_size: int, tp_size: int) -> int:
+    if moe_intermediate_size % tp_size != 0:
+        raise ValueError(
+            f"moe_intermediate_size ({moe_intermediate_size}) must be divisible by " f"tp_size ({tp_size})"
+        )
+
+    shard_size = moe_intermediate_size // tp_size
+    block_size = shard_size & (-shard_size)
+    block_size = min(block_size, 512)
+
+    return block_size
+
+
 def parse_quant_config(args, model_config, is_ernie, is_v1_loader):
     if args.quantization is not None and isinstance(args.quantization, str):
         args.quantization = parse_quantization(args.quantization)
@@ -89,7 +102,12 @@ def parse_quant_config(args, model_config, is_ernie, is_v1_loader):
             quantization_config["dense_quant_type"] = "block_wise_fp8"
             quantization_config["moe_quant_type"] = "w4afp8"
             tp_size = getattr(args, "tensor_parallel_size", 1)
-            quantization_config["hadamard_block_size"] = 512 // tp_size
+            moe_intermediate_size = getattr(model_config, "moe_intermediate_size", None)
+            if moe_intermediate_size is not None:
+                hadamard_block_size = _compute_hadamard_block_size(moe_intermediate_size, tp_size)
+                quantization_config["hadamard_block_size"] = hadamard_block_size
+            else:
+                quantization_config["hadamard_block_size"] = 512
             quantization_config["quantization"] = "mix_quant"
             quant_config_name = "mix_quant"
     else: