diff --git a/.gitignore b/.gitignore index b03a3562..e93eda4e 100644 --- a/.gitignore +++ b/.gitignore @@ -85,6 +85,7 @@ outputs/ running_logs/ .cursor/ _sglang/ +_tokenspeed/ wandb/log.txt .claude/ diff --git a/README.md b/README.md index 8329c825..cc3d945b 100644 --- a/README.md +++ b/README.md @@ -103,14 +103,21 @@ micromamba activate torchspec # Or install with SGLang ./tools/build_conda.sh micromamba activate torchspec + +# Or install TokenSpeed from an editable source checkout +./tools/build_conda.sh 1 tokenspeed +micromamba activate torchspec ``` To install into your current environment instead: ```bash -./tools/build_conda.sh current sglang # or 'vllm' or 'both' +./tools/build_conda.sh current tokenspeed # or 'sglang', 'vllm', or 'both' ``` +The TokenSpeed backend currently requires a Python 3.12 environment because +its native kernel dependency wheels do not support Python 3.14. + Optional: install Flash Attention support: ```bash diff --git a/configs/tokenspeed_qwen3_8b.yaml b/configs/tokenspeed_qwen3_8b.yaml new file mode 100644 index 00000000..9d4d5bce --- /dev/null +++ b/configs/tokenspeed_qwen3_8b.yaml @@ -0,0 +1,55 @@ +# Offline-first TokenSpeed integration smoke/validation config. +# +# Materialize: +# python -m torchspec.offline.generate \ +# --config configs/tokenspeed_qwen3_8b.yaml \ +# --output outputs/tokenspeed-qwen3-8b-offline-1000 + +model: + target_model_path: Qwen/Qwen3-8B + trust_remote_code: true + +dataset: + train_data_path: ../examples/data/sample_conversations.jsonl + chat_template: qwen + prompt_key: conversations + +training: + attention_backend: flex_attention + micro_batch_size: 1 + draft_accumulation_steps: 1 + max_concurrent_batches: 1 + max_seq_length: 2048 + num_epochs: 1 + training_num_gpus_per_node: 1 + training_num_nodes: 1 + ttt_length: 7 + +inference: + inference_engine_type: tokenspeed + inference_num_gpus: 1 + inference_num_gpus_per_engine: 1 + inference_num_gpus_per_node: 1 + inference_batch_size: 1 + inference_buffer_threshold: 8 + max_sample_pool_size: 16 + store_last_hidden_states: true + tokenspeed: + tp_size: 1 + nnodes: 1 + mem_fraction_static: 0.8 + init_timeout: 600 + +mooncake: + master_server_address: null + metadata_server: null + protocol: tcp + global_segment_size: 16GB + local_buffer_size: 4GB + +output_dir: ./outputs/tokenspeed-qwen3-8b +cache_dir: ./cache/tokenspeed-qwen3-8b +model_download_dir: /raid/hf/hub + +debug: + save_debug_train_data: null diff --git a/tools/build_conda.sh b/tools/build_conda.sh index fbe1787a..2ea87d05 100755 --- a/tools/build_conda.sh +++ b/tools/build_conda.sh @@ -12,18 +12,19 @@ PROJECT_ROOT="$(cd -- "$SCRIPT_DIR/.." && pwd)" # current - Install into current environment # 0 - Skip env creation and installation # BACKEND: -# sglang - Install SGLang only (default) -# vllm - Install vLLM only -# both - Install both backends +# sglang - Install SGLang only (default) +# vllm - Install vLLM only +# tokenspeed - Install TokenSpeed from an editable source checkout +# both - Install both backends MODE="${1:-1}" BACKEND="${2:-sglang}" # Validate backend -if [[ ! "$BACKEND" =~ ^(sglang|vllm|both)$ ]]; then +if [[ ! "$BACKEND" =~ ^(sglang|vllm|tokenspeed|both)$ ]]; then echo "Error: Invalid backend '$BACKEND'" echo "Usage: $0 [MODE] [BACKEND]" - echo " BACKEND options: sglang (default), vllm, both" + echo " BACKEND options: sglang (default), vllm, tokenspeed, both" exit 1 fi @@ -62,7 +63,7 @@ if [ "$MODE" = "1" ]; then "${ENV_CREATE_CMD[@]}" elif [ "$MODE" = "current" ]; then - echo "Using current environment: $(python3 --version), $(which python3)" + echo "Using current environment: $(python --version), $(command -v python)" else echo "Skipping environment setup (mode=0)" fi @@ -121,6 +122,75 @@ if [ "$BACKEND" = "vllm" ] || [ "$BACKEND" = "both" ]; then fi fi +# Install TokenSpeed if requested +if [ "$BACKEND" = "tokenspeed" ] && [ "$MODE" != "0" ]; then + echo "==========================================" + echo "Installing TokenSpeed..." + echo "==========================================" + + TOKENSPEED_REPO="${TOKENSPEED_REPO:-https://github.com/lightseekorg/tokenspeed.git}" + TOKENSPEED_FOLDER_NAME="${TOKENSPEED_FOLDER_NAME:-_tokenspeed}" + TOKENSPEED_PATH="${TOKENSPEED_PATH:-$PROJECT_ROOT/$TOKENSPEED_FOLDER_NAME}" + TOKENSPEED_REF="${TOKENSPEED_REF:-}" + + if [[ "$TOKENSPEED_PATH" != /* ]]; then + TOKENSPEED_PATH="$PROJECT_ROOT/$TOKENSPEED_PATH" + fi + + if [ -e "$TOKENSPEED_PATH" ] && [ ! -d "$TOKENSPEED_PATH/.git" ]; then + echo "Error: TOKENSPEED_PATH exists but is not a git checkout: $TOKENSPEED_PATH" + exit 1 + fi + + if [ ! -d "$TOKENSPEED_PATH/.git" ]; then + git clone "$TOKENSPEED_REPO" "$TOKENSPEED_PATH" + else + echo "Reusing existing TokenSpeed checkout: $TOKENSPEED_PATH" + fi + + if [ -n "$TOKENSPEED_REF" ]; then + echo "Checking out requested TokenSpeed ref: $TOKENSPEED_REF" + git -C "$TOKENSPEED_PATH" checkout "$TOKENSPEED_REF" + fi + + # TokenSpeed's native packages currently target Python 3.12. In particular, + # the kernel and CUDA dependency wheels do not resolve on Python 3.14. + TOKENSPEED_PYTHON_CHECK="import sys; assert sys.version_info[:2] == (3, 12), \ +f'TokenSpeed requires Python 3.12, got {sys.version.split()[0]}'" + + # Match TokenSpeed's published development-install instructions. The + # variable is needed by the runner image's system Python and is harmless in + # an isolated conda environment. + export PIP_BREAK_SYSTEM_PACKAGES=1 + + # Follow TokenSpeed's NVIDIA Docker build order. Installing the in-tree + # kernel first satisfies the runtime's tokenspeed-kernel>=0.1.3.dev0 + # dependency without trying to resolve an unavailable development wheel. + if [ "$MODE" = "1" ]; then + "${ENV_RUN_CMD[@]}" python -c "$TOKENSPEED_PYTHON_CHECK" + "${ENV_RUN_CMD[@]}" python -m pip install "setuptools==69.5.1" wheel + "${ENV_RUN_CMD[@]}" python -m pip install \ + -e "$TOKENSPEED_PATH/tokenspeed-kernel/python" \ + --no-build-isolation + "${ENV_RUN_CMD[@]}" python -m pip install \ + -e "$TOKENSPEED_PATH/tokenspeed-scheduler" + "${ENV_RUN_CMD[@]}" python -m pip install \ + -e "$TOKENSPEED_PATH/python" \ + --no-build-isolation + elif [ "$MODE" = "current" ]; then + python -c "$TOKENSPEED_PYTHON_CHECK" + python -m pip install "setuptools==69.5.1" wheel + python -m pip install \ + -e "$TOKENSPEED_PATH/tokenspeed-kernel/python" \ + --no-build-isolation + python -m pip install \ + -e "$TOKENSPEED_PATH/tokenspeed-scheduler" + python -m pip install \ + -e "$TOKENSPEED_PATH/python" \ + --no-build-isolation + fi +fi + # Install torchspec with appropriate extras if [ "$MODE" = "1" ]; then echo "==========================================" @@ -152,6 +222,9 @@ if [ "$MODE" = "1" ]; then echo "Backends: SGLang + vLLM" echo "SGLang: ./examples/qwen3-8b-single-node/run.sh" echo "vLLM: ./examples/qwen3-8b-single-node/run.sh --config configs/vllm_qwen3_8b.yaml" + elif [ "$BACKEND" = "tokenspeed" ]; then + echo "Backend: TokenSpeed" + echo "Source: $TOKENSPEED_PATH" fi elif [ "$MODE" = "current" ]; then EXTRAS="dev" @@ -181,5 +254,13 @@ else echo " pip install -e \"${SGLANG_FOLDER_NAME}/python[all]\"" echo " pip install vllm>=0.16.0" echo " pip install -e \".[dev,vllm]\"" + elif [ "$BACKEND" = "tokenspeed" ]; then + echo " git clone https://github.com/lightseekorg/tokenspeed.git _tokenspeed" + echo " export PIP_BREAK_SYSTEM_PACKAGES=1" + echo " pip install setuptools==69.5.1 wheel" + echo " pip install -e \"_tokenspeed/tokenspeed-kernel/python\" --no-build-isolation" + echo " pip install -e \"_tokenspeed/tokenspeed-scheduler\"" + echo " pip install -e \"_tokenspeed/python\" --no-build-isolation" + echo " pip install -e \".[dev]\"" fi fi diff --git a/torchspec/config/inference_config.py b/torchspec/config/inference_config.py index 0e3a7d32..3d08f521 100644 --- a/torchspec/config/inference_config.py +++ b/torchspec/config/inference_config.py @@ -133,6 +133,21 @@ class TrtllmConfig: extra_args: Dict[str, Any] = field(default_factory=dict) +@dataclass +class TokenSpeedConfig: + """TokenSpeed offline hidden-state inference configuration. + + The first integration targets single-node eager prefill. Tensor parallelism + is derived from ``inference_num_gpus_per_engine``. + """ + + tp_size: int = 1 + nnodes: int = 1 + mem_fraction_static: float = 0.8 + init_timeout: int = 600 + extra_args: Dict[str, Any] = field(default_factory=dict) + + @dataclass class OfflineTrainingConfig: """Configuration for training from materialized target outputs.""" @@ -160,6 +175,7 @@ class InferenceConfig: sglang: SGLangConfig = field(default_factory=SGLangConfig) vllm: VllmConfig = field(default_factory=VllmConfig) trtllm: TrtllmConfig = field(default_factory=TrtllmConfig) + tokenspeed: TokenSpeedConfig = field(default_factory=TokenSpeedConfig) def resolve_last_hidden_states_prenorm(self) -> bool: """Whether last_hidden_states from the engine are pre-norm. diff --git a/torchspec/config/train_config.py b/torchspec/config/train_config.py index 0995f393..8e160286 100644 --- a/torchspec/config/train_config.py +++ b/torchspec/config/train_config.py @@ -329,6 +329,7 @@ def load_config( "sglang": "sglang_", "vllm": "vllm_", "trtllm": "trtllm_", + "tokenspeed": "tokenspeed_", } diff --git a/torchspec/inference/engine/tokenspeed_engine.py b/torchspec/inference/engine/tokenspeed_engine.py new file mode 100644 index 00000000..a4ed4f60 --- /dev/null +++ b/torchspec/inference/engine/tokenspeed_engine.py @@ -0,0 +1,288 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Ray actor wrapper for TokenSpeed offline hidden-state materialization.""" + +from __future__ import annotations + +from typing import Any + +import ray +import torch +from omegaconf import DictConfig, OmegaConf + +from torchspec.inference.engine.base import InferenceEngine +from torchspec.ray.ray_actor import RayActor +from torchspec.transfer.mooncake.eagle_store import HIDDEN_STATES_STORAGE_DTYPE +from torchspec.utils.logging import logger, setup_file_logging + +_PROTECTED_ENGINE_KEYS = frozenset( + { + "model", + "attn_tp_size", + "base_gpu_id", + "nprocs_per_node", + "nnodes", + "gpu_memory_utilization", + "enable_spec_training_mooncake", + "eagle3_layers_to_capture", + "enable_prefix_caching", + "disable_kvstore", + "enforce_eager", + "disable_overlap_schedule", + "disable_prefill_graph", + "skip_server_warmup", + "chunked_prefill_size", + "max_model_len", + "max_prefill_tokens", + } +) + + +class TokenSpeedEngine(InferenceEngine, RayActor): + """Use TokenSpeed eager prefill to publish TorchSpec Mooncake records.""" + + def __init__( + self, + args, + rank: int, + base_gpu_id: int | None = None, + num_gpus_per_engine: int = 1, + node_rank: int = 0, + engine_group: int = 0, + ) -> None: + self.args = args + self.rank = rank + self.base_gpu_id = base_gpu_id + self.num_gpus_per_engine = num_gpus_per_engine + self.node_rank = node_rank + self.local_gpu_id = None + self._engine = None + self._mooncake_config = None + self._hidden_size = None + self.aux_hidden_state_layer_ids: list[int] = [] + setup_file_logging("inference", self.rank, group=engine_group) + + def init( + self, + mooncake_config=None, + dist_init_addr: str | None = None, + pre_allocated_port: int | None = None, + ) -> None: + del dist_init_addr, pre_allocated_port + if self.node_rank != 0 or getattr(self.args, "tokenspeed_nnodes", 1) != 1: + raise NotImplementedError("The initial TokenSpeed integration supports one node only") + if not getattr(self.args, "store_last_hidden_states", True): + raise NotImplementedError( + "TokenSpeed currently requires inference.store_last_hidden_states=true" + ) + if getattr(self.args, "train_with_decode", False): + raise NotImplementedError( + "TokenSpeed currently supports offline prefill, not train_with_decode" + ) + + if self.base_gpu_id is not None: + self.local_gpu_id = self.setup_gpu(self.base_gpu_id) + else: + self.local_gpu_id = self.setup_gpu() + + self._mooncake_config = mooncake_config + if mooncake_config is None: + raise ValueError("TokenSpeed hidden-state capture requires Mooncake") + mooncake_config.local_hostname = self.get_node_ip() + mooncake_config.export_env() + + from torchspec.transfer.mooncake.utils import ( + check_mooncake_master_available, + ) + + check_mooncake_master_available( + mooncake_config.master_server_address, + mooncake_config.metadata_server, + ) + + from transformers import AutoConfig + + model_config = AutoConfig.from_pretrained( + self.args.target_model_path, + trust_remote_code=getattr(self.args, "trust_remote_code", True), + cache_dir=getattr(self.args, "model_download_dir", None), + ) + model_config = getattr(model_config, "text_config", model_config) + self._hidden_size = int(model_config.hidden_size) + if self.args.aux_hidden_states_layers is not None: + self.aux_hidden_state_layer_ids = list(self.args.aux_hidden_states_layers) + else: + num_layers = int(model_config.num_hidden_layers) + self.aux_hidden_state_layer_ids = [ + 1, + num_layers // 2 - 1, + num_layers - 4, + ] + + tp_size = self.num_gpus_per_engine + configured_tp = getattr(self.args, "tokenspeed_tp_size", tp_size) + if configured_tp != tp_size: + raise ValueError( + f"tokenspeed.tp_size ({configured_tp}) must equal " + f"inference_num_gpus_per_engine ({tp_size})" + ) + + extra_args = getattr(self.args, "tokenspeed_extra_args", None) + if isinstance(extra_args, DictConfig): + extra = OmegaConf.to_container(extra_args, resolve=True) + else: + extra = dict(extra_args or {}) + blocked = extra.keys() & _PROTECTED_ENGINE_KEYS + if blocked: + logger.warning( + "TokenSpeed extra_args contains managed keys that will be ignored: %s", + sorted(blocked), + ) + extra = {key: value for key, value in extra.items() if key not in blocked} + + max_seq_length = int(getattr(self.args, "max_seq_length", 8192)) + engine_kwargs = { + "log_level": "warning", + **extra, + "model": self.args.target_model_path, + "attn_tp_size": tp_size, + "base_gpu_id": self.local_gpu_id, + "nprocs_per_node": tp_size, + "nnodes": 1, + "gpu_memory_utilization": getattr(self.args, "tokenspeed_mem_fraction_static", 0.8), + "enable_spec_training_mooncake": True, + "eagle3_layers_to_capture": ",".join( + str(layer_id) for layer_id in self.aux_hidden_state_layer_ids + ), + "max_model_len": max_seq_length, + "max_prefill_tokens": max_seq_length, + "trust_remote_code": getattr(self.args, "trust_remote_code", True), + "download_dir": getattr(self.args, "model_download_dir", None), + } + + from tokenspeed.runtime.entrypoints.engine import Engine + + self._engine = Engine(**engine_kwargs) + logger.info( + "TokenSpeedEngine rank %d initialized: model=%s tp=%d aux_layers=%s", + self.rank, + self.args.target_model_path, + tp_size, + self.aux_hidden_state_layer_ids, + ) + + def generate( + self, + data_id: str | list[str], + input_ids_ref: ray.ObjectRef | list[torch.Tensor] | None = None, + packed_loss_mask_list: list[str] | None = None, + formatted_prompts: list[str] | None = None, + return_last_hidden_states: bool = False, + return_logits: bool = True, + multimodal_inputs: list[dict] | None = None, + ) -> list[dict[str, Any]]: + del data_id, packed_loss_mask_list, return_last_hidden_states, return_logits + if self._engine is None: + raise RuntimeError("TokenSpeedEngine is not initialized") + if any(bool(item) for item in (multimodal_inputs or [])): + raise NotImplementedError( + "The initial TokenSpeed integration does not support multimodal inputs" + ) + if (input_ids_ref is None) == (formatted_prompts is None): + raise ValueError("Exactly one of input_ids_ref or formatted_prompts must be set") + + input_ids_list: list[list[int]] | None = None + if formatted_prompts is not None: + request_kwargs = {"prompt": formatted_prompts} + expected_lengths = None + else: + resolved = ( + ray.get(input_ids_ref) + if isinstance(input_ids_ref, ray.ObjectRef) + else input_ids_ref + ) + input_ids_list = [] + for ids in resolved: + if ids.dim() == 2 and ids.shape[0] == 1: + ids = ids.squeeze(0) + if ids.dim() != 1: + raise ValueError(f"Unexpected input_ids shape: {ids.shape}") + input_ids_list.append(ids.tolist()) + request_kwargs = {"input_ids": input_ids_list} + expected_lengths = [len(ids) for ids in input_ids_list] + + results = self._engine.generate( + **request_kwargs, + sampling_params={"max_new_tokens": 0, "temperature": 0}, + return_hidden_states=True, + ) + if isinstance(results, dict): + results = [results] + + outputs = [] + for index, result in enumerate(results): + keys = result.get("meta_info", {}).get("spec_training_mooncake_store_keys", []) + if len(keys) != 1: + raise RuntimeError( + f"TokenSpeed did not return exactly one Mooncake key for result {index}: {keys}" + ) + seq_len = result["meta_info"].get("prompt_tokens") + if seq_len is None and expected_lengths is not None: + seq_len = expected_lengths[index] + if seq_len is None: + raise RuntimeError("TokenSpeed did not report prompt_tokens") + outputs.append( + { + "mooncake_key": keys[0], + "tensor_shapes": { + "hidden_states": ( + seq_len, + len(self.aux_hidden_state_layer_ids) * self._hidden_size, + ), + "input_ids": (seq_len,), + "last_hidden_states": (seq_len, self._hidden_size), + }, + "tensor_dtypes": { + "hidden_states": HIDDEN_STATES_STORAGE_DTYPE, + "input_ids": torch.long, + "last_hidden_states": HIDDEN_STATES_STORAGE_DTYPE, + }, + } + ) + return outputs + + def health_check(self, timeout: float = 5.0) -> bool: + del timeout + return self._engine is not None + + def shutdown(self) -> None: + if self._engine is not None: + self._engine.shutdown() + self._engine = None + logger.info("TokenSpeedEngine rank %d shutdown complete", self.rank) + + def get_status(self) -> dict: + return { + "rank": self.rank, + "initialized": self._engine is not None, + "base_gpu_id": self.base_gpu_id, + "hidden_size": self._hidden_size, + } diff --git a/torchspec/inference/factory.py b/torchspec/inference/factory.py index dc2f52c9..650c922b 100644 --- a/torchspec/inference/factory.py +++ b/torchspec/inference/factory.py @@ -46,7 +46,7 @@ def create_inference_engines(args, inference_pg, mooncake_config, engine_group: _wait_for_init(init_refs, "OfflineReplay", timeout=300) return engines - if engine_type not in ("hf", "sgl", "vllm", "trtllm"): + if engine_type not in ("hf", "sgl", "vllm", "trtllm", "tokenspeed"): raise ValueError(f"Unknown inference_engine_type: {engine_type}") logger.info(f"Using {engine_type} engine for inference") @@ -80,7 +80,7 @@ def prepare_inference_engines(args, inference_pg, mooncake_config, engine_group: if engine_type == "offline": return _prepare_offline_replay_engines(args, mooncake_config, engine_group) - if engine_type not in ("hf", "sgl", "vllm", "trtllm"): + if engine_type not in ("hf", "sgl", "vllm", "trtllm", "tokenspeed"): raise ValueError(f"Unknown inference_engine_type: {engine_type}") logger.info(f"Preparing {engine_type} inference engines...") @@ -93,6 +93,10 @@ def prepare_inference_engines(args, inference_pg, mooncake_config, engine_group: engines, init_refs = _prepare_trtllm_engines( args, inference_pg, mooncake_config, engine_group ) + elif engine_type == "tokenspeed": + engines, init_refs = _prepare_tokenspeed_engines( + args, inference_pg, mooncake_config, engine_group + ) else: engines, init_refs = _prepare_vllm_engines( args, inference_pg, mooncake_config, engine_group @@ -145,6 +149,8 @@ def init_engines(args, pg, engine_type: str, mooncake_config=None, engine_group: return _init_vllm_engines(args, pg, mooncake_config, engine_group) elif engine_type == "trtllm": return _init_trtllm_engines(args, pg, mooncake_config, engine_group) + elif engine_type == "tokenspeed": + return _init_tokenspeed_engines(args, pg, mooncake_config, engine_group) else: raise ValueError(f"Unknown engine_type: {engine_type}") @@ -547,6 +553,68 @@ def _init_trtllm_engines(args, pg, mooncake_config=None, engine_group: int = 0) return engines +def _prepare_tokenspeed_engines( + args, pg, mooncake_config=None, engine_group: int = 0 +) -> tuple[list, list]: + """Create single-node TokenSpeed engine actors.""" + nnodes = getattr(args, "tokenspeed_nnodes", 1) + if nnodes != 1: + raise NotImplementedError( + "The initial TokenSpeed backend supports single-node inference only" + ) + + num_gpus_total = getattr(args, "inference_num_gpus", 1) + gpus_per_engine = getattr(args, "inference_num_gpus_per_engine", 1) + num_engines = num_gpus_total // gpus_per_engine + if num_engines <= 0: + raise ValueError("TokenSpeed requires at least one inference engine") + + from torchspec.inference.engine.tokenspeed_engine import TokenSpeedEngine + + pg_obj, reordered_bundle_indices, reordered_gpu_ids = pg + TokenSpeedRayActor = ray.remote(TokenSpeedEngine) + env_vars = get_torchspec_env_vars() + engines = [] + for index in range(num_engines): + bundle_offset = index * gpus_per_engine + base_gpu_id = int(reordered_gpu_ids[bundle_offset]) + scheduling_strategy = PlacementGroupSchedulingStrategy( + placement_group=pg_obj, + placement_group_capture_child_tasks=True, + placement_group_bundle_index=reordered_bundle_indices[bundle_offset], + ) + engines.append( + TokenSpeedRayActor.options( + num_cpus=0.2, + num_gpus=0.2, + scheduling_strategy=scheduling_strategy, + runtime_env={"env_vars": env_vars}, + ).remote( + args=args, + rank=index, + base_gpu_id=base_gpu_id, + num_gpus_per_engine=gpus_per_engine, + node_rank=0, + engine_group=engine_group, + ) + ) + + init_refs = [engine.init.remote(mooncake_config=mooncake_config) for engine in engines] + logger.info( + "Preparing %d TokenSpeed engine(s), %d GPU(s) each", + num_engines, + gpus_per_engine, + ) + return engines, init_refs + + +def _init_tokenspeed_engines(args, pg, mooncake_config=None, engine_group: int = 0) -> list: + engines, init_refs = _prepare_tokenspeed_engines(args, pg, mooncake_config, engine_group) + timeout = getattr(args, "tokenspeed_init_timeout", 600) + _wait_for_init(init_refs, "TokenSpeed", timeout=timeout) + return engines + + def _create_and_init_actors( args, pg,