feat: tokenspeed support

.gitignore

@@ -85,6 +85,7 @@ outputs/
 running_logs/
 .cursor/
 _sglang/
+_tokenspeed/
 wandb/log.txt
 
 .claude/

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

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

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

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.

torchspec/config/train_config.py

@@ -329,6 +329,7 @@ def load_config(
     "sglang": "sglang_",
     "vllm": "vllm_",
     "trtllm": "trtllm_",
+    "tokenspeed": "tokenspeed_",
 }