FastTTS: A Serving Framework for Test-Time Scaling on the Edge

Paper | Code | Project Page

FastTTS is a high-performance Test-Time Scaling (TTS) framework designed for memory-constrained edge language models. Built as a plug-and-play library on top of vLLM, it pairs a generator LLM with a verifier (e.g., a Process Reward Model) and orchestrates verifier-guided search strategies — such as beam search, diverse tree search (DVTS), and best-of-N — to find high-quality reasoning traces at inference time. By eliminating system-level bottlenecks, FastTTS enables edge LLMs on single consumer GPUs to match the accuracy and latency of large cloud models.

Key design choices:

• System-level optimizations: As introduced in our paper, the framework integrates (i) Speculative Beam Extension to mitigate system stragglers caused by irregular reasoning paths, (ii) Asymmetric Multi-Model Memory Allocation to intelligently partition memory between the generator and verifier, and (iii) Dynamic Prefix-Aware Scheduling to maximize KV cache reuse from dynamic prefix sharing.
• Separate processes for the generator and verifier, enabling offloading and independent GPU memory management.
• Configurable per-request search parameters via SearchConfig.

Todo

• Add the asymmetric multi-model memory allocation implementation
• Add benchmark utility
• Add intelligent speculative candidate selection

---

Repository Structure

FastTTS-os/
├── __init__.py                  # Package exports
├── config.py                    # FastTTSConfig & SearchConfig dataclasses
├── fasttts.py                   # FastTTS class and create_fasttts() entry point
├── setup.py                     # Package installation
├── environment.yml              # Conda environment specification
│
├── models/
│   ├── __init__.py
│   ├── vllm_wrapper.py          # Generator / Verifier wrappers (multi-process)
│   ├── tts_llm.py               # TTSLLM — custom vLLM LLM subclass
│   ├── generator_engine.py      # Generator LLM engine with spec beam extension
│   ├── verifier_engine.py       # Verifier LLM engine
│   ├── custom_scheduler.py      # Custom vLLM scheduler
│   ├── spec_stopchecker.py      # Stop checker for speculative beam extension
│   ├── reward_utils.py          # Reward model input preparation & sigmoid
│   └── numbers.py               # Priority constants
│
└── search/
    ├── __init__.py
    ├── beam.py                  # Beam dataclass
    ├── beam_search.py           # Beam search implementation
    ├── dvts.py                  # Diverse Tree Search (DVTS) implementation
    ├── best_of_n.py             # Best-of-N sampling implementation
    └── utils.py                 # Shared utilities (conversation builder, scoring, truncation)

---

Installation

git clone https://github.com/ihc-fan-lab/FastTTS
cd FastTTS
pip install -e .

Requirements

• Python >= 3.10
• PyTorch >= 2.4
• vLLM >= 0.6
• transformers >= 4.45

---

Quick Start

# 1. Load a dataset
dataset = load_dataset("HuggingFaceH4/aime_2024", split="train")
problems = [dataset[0]["problem"]]

# 2. Create a FastTTS instance
fasttts = create_fasttts(
    generator_vllm_config={
        "model": "Qwen/Qwen2.5-Math-1.5B-Instruct",
        "gpu_memory_utilization": 0.45,
    },
    verifier_vllm_config={
        "model": "Skywork/Skywork-o1-Open-PRM-Qwen-2.5-1.5B",
        "gpu_memory_utilization": 0.45,
    },
    approach="beam_search",
    offload_enabled=False, 
    spec_beam_extension=True, 
    prefix_aware_scheduling=True,
)

# 3. Configure the search
search_config = SearchConfig(
    approach="beam_search",   # "beam_search", "dvts", or "best_of_n"
    beam_width=4,
    n=8,
    num_iterations=3,
    temperature=0.8,
)

# 4. Run
try:
    fasttts.initialize()
    results = fasttts.search(problems, search_config=search_config)
    avg_completion_time = sum(results['completion_time'][0]) / (len(results['completion_time'][0]))
    print(f"Average completion time: {avg_completion_time:.2f}s")
    acg_precise_goodput = sum(results['effective_num_tokens'][0]) / (len(results['effective_num_tokens'][0]))
    print(f"Average precise goodput: {acg_precise_goodput:.2f}")
finally:
    fasttts.shutdown()

---

Search Approaches

Beam Search (beam_search)

Standard beam search with PRM-guided pruning. At each step, n beams are generated, scored by the verifier, and the top n / beam_width are kept for expansion.

Diverse Tree Search (dvts)

Maintains independent subtrees to promote diversity. Each subtree selects its own best beam, preventing the search from collapsing to a single reasoning path.

Best-of-N (best_of_n)

Generates n independent completions in a single forward pass and selects the highest-scoring one via the verifier. Simplest strategy with the lowest latency overhead.

---

API Reference

create_fasttts(...) -> FastTTS

Factory function to create a FastTTS instance.

fasttts = create_fasttts(
    generator_vllm_config={"model": "...", ...},
    verifier_vllm_config={"model": "...", ...},
    approach="beam_search",
    # Any FastTTSConfig or SearchConfig field can be passed as a kwarg
    offload_enabled=False,
    spec_beam_extension=False,
    prefix_aware_scheduling=False,
)

FastTTS

Method	Description
initialize()	Load generator and verifier models. Called automatically on first search() if needed.
search(problems, search_config=None, **kwargs)	Run search over a list of problem strings. Returns a results dict.
search_single(problem, search_config=None, **kwargs)	Convenience wrapper for a single problem.
create_search_config(**kwargs)	Create a SearchConfig from the instance defaults with overrides.
shutdown()	Release model resources.

FastTTS also supports use as a context manager:

with create_fasttts(...) as fasttts:
    fasttts.initialize()
    results = fasttts.search(problems, search_config=config)

Results Dictionary

The dict returned by search() contains:

Key	Type	Description
completions	list[list[str]]	All generated completions per problem
pred	list[str]	Best completion per problem (highest aggregated score)
scores	list[list[list[float]]]	Per-step PRM scores for each completion
completion_tokens	list[list[int]]	Token counts per completion
effective_num_tokens	list[list[int]]	Total tokens generated per beam (including extended)
total_num_tokens	int	Total tokens generated across the entire search
n_completion_tokens	int	Tokens generated to collect n completions
total_generator_latency_s	float	Total wall-clock generator time (seconds)
total_verifier_latency_s	float	Total wall-clock verifier time (seconds)
n_generator_latency_s	float	Generator time to collect n completions
n_verifier_latency_s	float	Verifier time to collect n completions
completion_time	list[list[float]]	Time at which each beam completed
extended_tokens_list	list	Tokens from speculative beam extension per iteration

---

Configuration Reference

FastTTSConfig

Top-level configuration passed to the FastTTS constructor.

Parameter	Type	Default	Description
generator_vllm_config	dict	(see below)	vLLM engine kwargs for the generator model
verifier_vllm_config	dict	(see below)	vLLM engine kwargs for the verifier model
temperature	float	0.8	Default sampling temperature
top_p	float	1.0	Default nucleus sampling threshold
max_tokens	int	2048	Default max generation length
system_prompt	str	(math prompt)	System prompt prepended to all conversations
custom_chat_template	str | None	None	Override the tokenizer chat template
search_config	SearchConfig	SearchConfig()	Default search configuration
spec_beam_extension	bool	False	Enable speculative beam extension
offload_enabled	bool	False	Enable model sleep/wake for memory savings
prefix_aware_scheduling	bool	False	Enable prefix-aware request scheduling

Default generator_vllm_config:

{
    "model": "Qwen/Qwen2.5-Math-1.5B-Instruct",
    "max_model_len": 4096,
    "gpu_memory_utilization": 0.45,
    "tensor_parallel_size": 1,
    "enable_prefix_caching": True,
    "seed": 42,
    "enable_chunked_prefill": False,
    "disable_log_stats": False,
}

Default verifier_vllm_config:

{
    "model": "Skywork/Skywork-o1-Open-PRM-Qwen-2.5-1.5B",
    "max_model_len": 4096,
    "gpu_memory_utilization": 0.45,
    "tensor_parallel_size": 1,
    "enable_prefix_caching": True,
    "seed": 42,
    "disable_log_stats": False,
}

SearchConfig

Per-request search parameters. Create via SearchConfig(...) or fasttts.create_search_config(...).

Parameter	Type	Default	Description
Search approach
approach	str	"beam_search"	"beam_search", "dvts", or "best_of_n"
Generation
temperature	float	0.8	Sampling temperature
top_p	float	1.0	Nucleus sampling threshold
max_tokens	int	2048	Maximum tokens per generation step
stop	str | None	"\n\n"	Stop string for step-level generation
Beam search / DVTS
beam_width	int	4	Expansion factor per beam (kept beams = n / beam_width)
num_iterations	int	40	Maximum search iterations
lookahead	int	0	Greedy lookahead steps for scoring (0 = none)
n	int	8	Total number of parallel beams (must be a multiple of beam_width)
Best-of-N
num_samples	int	4	Number of samples to generate
Scoring
agg_strategy	str	"last"	Score aggregation: "last", "min", "prod", or "mean"
filter_duplicates	bool	False	Remove duplicate completions
sort_completed	bool	False	Sort completed beams by score
Truncation (speculative beam extension)
truncation_mean_ratio	float	0.85	Mean ratio of tokens to keep when truncating duplicate beams
truncation_std_ratio	float	0.1	Standard deviation of the truncation ratio
truncation_min_tokens	int	1	Minimum number of tokens to keep after truncation
Prompt
system_prompt	str	(math prompt)	System prompt for the conversation
custom_chat_template	str | None	None	Override the tokenizer chat template
Batch processing
batch_size	int	1	Number of problems per batch (must be 1 for beam_search)

Supported Verifier Models

Model	Hub ID
Skywork PRM 1.5B	Skywork/Skywork-o1-Open-PRM-Qwen-2.5-1.5B
Skywork PRM 7B	Skywork/Skywork-o1-Open-PRM-Qwen-2.5-7B
Math-Shepherd PRM	peiyi9979/math-shepherd-mistral-7b-prm

---

Citation

@article{chen2025fasttts,
  title={FastTTS: Accelerating Test-Time Scaling for Edge LLM Reasoning},
  author={Chen, Hao Mark and Mo, Zhiwen and Lu, Guanxi and Liang, Shuang and Ma, Lingxiao and Luk, Wayne and Fan, Hongxiang},
  journal={arXiv preprint arXiv:2509.00195},
  year={2025}
}

Acknowledgment

This project adapts some parts of the code from Hugging Face Search and Learn. We thank the Hugging Face team for their contribution.