MEGA Hub

PyroDash: Cost-Efficient Token-Level Small-Large Language Model Collaborative Inference

Authors

Do you know Niqi Lyu?You can claim authorship or link another user.Do you know Pengtao Shi?You can claim authorship or link another user.Do you know Wei Qiu?You can claim authorship or link another user.Do you know Jianlin Zhong?You can claim authorship or link another user.Do you know Sicong Xia?You can claim authorship or link another user.Do you know Jianyao Ma?You can claim authorship or link another user.Do you know Yicheng Ding?You can claim authorship or link another user.

Abstract

Large language models (LLMs) provide strong reasoning capabilities but are expensive to serve at scale, whereas small language models (SLMs) are cheaper but less reliable on difficult problems. We introduce PyroDash, a cost-aware framework for token-level SLM-LLM collaborative inference. During generation, the SLM decides whether to request assistance by emitting a control token. A Collaborate Engine then sends the query and partial reasoning trace to a frozen LLM for completion through a single handoff. The policy is internalized in the SLM, requiring neither a separate router, LLM retraining, nor access to LLM logits. PyroDash trains the SLM in three stages: control-token embedding learning, offloading-oriented supervised fine-tuning, and cost-aware alignment with Group Relative Policy Optimization. Its reward balances answer accuracy against inference cost normalized by LLM-only inference. Across five mathematical reasoning benchmarks, PyroDash supports different accuracy-cost operating points. With $λ=0.05$, it achieves 64.04 percent average accuracy, 6.36 percentage points above the LLM-only baseline, while reducing cost by 20.4 percent. With $λ=0.6$, it achieves 54.55 percent accuracy with a 1.90 percent LLM token ratio and 0.012 LLM calls per example, reducing total cost from USD 49.36 to USD 1.78. These results show that learned token-level handoffs can reduce LLM use while preserving strong reasoning performance.

Community

00

Publication notes

Author note
19 pages, 3 figures