chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/runtime/managers/cpu_worker.py

@@ -0,0 +1,76 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import torch
+
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import (
+    init_logger,
+)
+from sglang.srt.utils import cpu_has_amx_support, get_cpu_ids_by_node
+
+from .gpu_worker import GPUWorker
+
+_is_cpu_amx_available = cpu_has_amx_support()
+
+logger = init_logger(__name__)
+
+
+class CPUWorker(GPUWorker):
+    """
+    A worker that executes the model on pure CPU platforms
+    """
+
+    def __init__(
+        self,
+        local_rank: int,
+        rank: int,
+        master_port: int,
+        server_args: ServerArgs,
+    ):
+        super().__init__(local_rank, rank, master_port, server_args)
+        if _is_cpu_amx_available:
+            self.init_cpu_threads_binding()
+
+    def init_cpu_threads_binding(self):
+        omp_cpuids = os.environ.get("SGLANG_CPU_OMP_THREADS_BIND", "all")
+        cpu_ids_by_node = get_cpu_ids_by_node()
+        n_numa_node = len(cpu_ids_by_node)
+        if omp_cpuids == "all":
+            assert self.server_args.tp_size <= n_numa_node, (
+                f"SGLANG_CPU_OMP_THREADS_BIND is not set, in this case, "
+                f"tp_size {self.server_args.tp_size} should be smaller than or equal to number of numa node on the machine {n_numa_node}. "
+                f"If you need tp_size to be larger than number of numa node, please set the CPU cores for each tp rank via SGLANG_CPU_OMP_THREADS_BIND explicitly. "
+                f"For example, on a machine with 2 numa nodes, where core 0-31 are on numa node 0 and core 32-63 are on numa node 1, "
+                f"it is suggested to use -tp 2 and bind tp rank 0 to core 0-31 and tp rank 1 to core 32-63. "
+                f"This is the default behavior if SGLANG_CPU_OMP_THREADS_BIND is not set and it is the same as setting SGLANG_CPU_OMP_THREADS_BIND=0-31|32-63. "
+                f"If you do need tp_size to be larger than the number of numa nodes, you could set SGLANG_CPU_OMP_THREADS_BIND explicitly for example SGLANG_CPU_OMP_THREADS_BIND=0-15|16-31|32-47|48-63 and run with -tp 4. "
+                f"If you don't want each tp rank to use all the cores on one numa node, you could set for example SGLANG_CPU_OMP_THREADS_BIND=0-15|32-47 and run with -tp 2."
+            )
+            if self.server_args.tp_size < n_numa_node:
+                logger.warning(
+                    f"Detected the current machine has {n_numa_node} numa nodes available, but tp_size is set to {self.server_args.tp_size}, so only {self.server_args.tp_size} numa nodes are used."
+                )
+            self.local_omp_cpuid = cpu_ids_by_node[self.rank]
+        else:
+            threads_bind_list = omp_cpuids.split("|")
+            assert self.server_args.tp_size == len(threads_bind_list), (
+                f"SGLANG_CPU_OMP_THREADS_BIND setting must be aligned with TP size parameter ({self.server_args.tp_size}). "
+                f"Please double check your settings."
+            )
+            self.local_omp_cpuid = threads_bind_list[self.rank]
+            if self.server_args.tp_size > n_numa_node:
+                logger.warning(
+                    f"TP size ({self.server_args.tp_size})is larger than numa node number ({n_numa_node}), "
+                    f"in this case the available memory amount of each rank cannot be determined in prior. "
+                    f"Please set proper `--max-total-tokens` to avoid the out-of-memory error."
+                )
+
+        # Bind OpenMP threads to CPU cores
+        torch.ops.sgl_kernel.init_cpu_threads_env(self.local_omp_cpuid)
+
+        # Set local size to hint SGLang to use shared memory based AllReduce
+        os.environ["LOCAL_SIZE"] = str(self.server_args.tp_size)
+        torch.ops.sgl_kernel.initialize(self.server_args.tp_size, self.rank)