Skip to content

Latest commit

 

History

2 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

RAPID

RAPID is a minimal reference implementation of the control mechanisms introduced in the RAPID paper. It isolates the parts needed to study adaptive draft-depth selection and interruptible proactive drafting without requiring a model runtime or GPU.

The repository contains:

  • CAAC, which combines recent token acceptance with measured local generation latency to select the next draft depth.
  • IPDG, which checks cloud feedback only after a complete expansion step. A root match keeps the proactive branch. A root mismatch stops before the next expansion and restores the saved tree and KV-cache boundary.
  • A round-level simulator that connects CAAC and IPDG through a small backend protocol.
  • CPU-only examples and unit tests for the controller and stopping behavior.

The project page provides a replayable IPDG timeline. It shows why expansion continues after a root match and stops at the next expansion boundary after a root mismatch.

Quick start

RAPID requires Python 3.10 or later. Its runtime has no third-party dependencies.

git clone https://github.com/jianpingpei/RAPID.git
cd RAPID
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -e ".[dev]"
python examples/simulate_round.py
python -m pytest

The example uses an in-memory backend. It does not load model weights or contact a cloud server.

Python API

The public API is intentionally small.

API Role
CAACConfig Stores the controller parameters and depth bounds.
CAACController.update(current_acceptance, per_token_latency) Updates the controller state and returns a CAACDecision.
run_ipdg_expansion(...) Runs boundary-aligned proactive expansion and returns an IPDGExpansionResult.
RapidRoundOps Defines the operations supplied by a real tree and KV-cache backend.
run_round(...) Runs one RAPID round and returns a RapidRoundResult.

The implementation is split across three modules:

src/rapid/caac.py    CAAC state and update rule
src/rapid/ipdg.py    IPDG expansion and stopping rule
src/rapid/round.py   round orchestration and backend protocol

Paper parameters

The defaults follow the configuration used for the paper's main results.

Parameter Default Purpose
Target acceptance, A* 0.75 Reference for the logical path
Initial acceptance EMA 0.8 Bootstrap value before the first verification result
Kp, Ki, Kd 5, 0.1, 1 PID gains
Acceptance EMA weight, beta 0.6 Weight of the latest acceptance observation
Low-acceptance reset 0.6 Clears the stored integral state after the current update
Fast-latency weight, gamma 0.5 Weight of the latest latency observation
Baseline adaptation rate, kappa 0.01 Rate at which the slow latency baseline rises
Slowdown strength, alpha 0.5 Maps ESR to the physical-path penalty
ESR cap 0.99 Keeps the slowdown calculation bounded
Integral-state range [-1, 1] Limits accumulated PID error
Draft-depth range 1 to 8 Bounds the planned expansion depth

CAAC keeps a continuous depth state and quantizes the selected depth with ceil. IPDG can stop before that planned depth when cloud verification reports a root mismatch.

Integrating a serving backend

RapidRoundOps separates the control logic from a particular inference stack. A production adapter must provide operations that:

  1. save a recoverable draft-tree and KV-cache boundary;
  2. send the current draft for asynchronous verification;
  3. start and expand the proactive branch one complete step at a time;
  4. check whether feedback is ready and whether the proactive root matches;
  5. retain and remap a valid branch, or restore and discard an invalid branch; and
  6. expose the acceptance and local-latency observations used by CAAC.

The expansion call must leave the tree and KV cache at a consistent boundary before IPDG checks feedback. If feedback arrives while a step is running, that step finishes before IPDG decides whether to continue. If a backend callback raises, the surrounding request handler remains responsible for cancelling the cloud request and restoring its saved boundary.

Repository boundary

This repository is a paper-facing reference implementation. It is not a complete artifact or a serving deployment. It does not include:

  • model weights or datasets;
  • raw measurements or reported result files;
  • experiment orchestration, workload injection, or hardware-monitoring scripts;
  • machine addresses, hardware-specific configuration, or deployment credentials; or
  • the complete SpecEdge serving stack.

These omissions keep the public code focused on the two RAPID mechanisms. They also mean that this repository alone cannot reproduce the paper's end-to-end performance tables.

Paper and citation

Paper: Coming soon.

Citation metadata will be added after publication. Until then, please link to this repository when referring to the implementation.

License and upstream work

RAPID is based on SpecEdge. The included license permits research, evaluation, and educational use, but it is not the MIT License. Direct commercial deployment, productization, or sale requires a separate commercial license from SpecEdge. See LICENSE and NOTICE for the full terms and source revision.

About

Minimal reference implementation of RAPID edge-cloud speculative decoding controls

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages