pipealloc

Token budget allocation optimizer for multi-step LLM pipelines.

Problem

Multi-step LLM pipelines (retrieval, summarization, generation, validation) operate under a fixed total token budget. The standard approach is to split that budget equally across steps. This is wasteful: some steps saturate quickly (a validation check needs only 200 tokens to be effective), while others are token-hungry (chain-of-thought generation improves continuously up to thousands of tokens). Uniform allocation starves the steps that need tokens most and wastes tokens on steps that stopped benefiting long ago.

There is no tool that models this tradeoff and solves for the optimal split. Developers tune allocations by hand, guessing at diminishing returns they cannot see.

What It Does

pipealloc models each pipeline step's quality-vs-tokens curve, accounts for how upstream errors propagate to downstream steps, and uses constrained optimization to find the allocation that maximizes end-to-end pipeline quality for any given total budget.

Given a pipeline definition and a token budget, it returns:

• The optimal token allocation per step
• The expected quality improvement over uniform allocation
• A sensitivity analysis showing which steps are bottlenecks

Why It Is Interesting

1. The optimization is non-trivial. Error propagation between steps creates coupling: investing in an early step benefits all downstream steps, not just itself. This makes the problem harder than independent knapsack allocation.

2. The gains are large. On a realistic 4-step RAG pipeline, pipealloc achieves 7 to 18% quality improvement over uniform allocation at the same total token cost (see Results below). At tight budgets, the difference between a working pipeline and a broken one can come down to allocation.

3. Water-filling beats gradient descent. Standard optimizers (SLSQP) get trapped in local minima when step curves have very different saturation rates. pipealloc uses a multi-resolution water-filling algorithm (allocate in chunks of 50, 10, then 1 token, always to the step with highest marginal gain) followed by a gradient-based polish. This reliably finds the global optimum.

4. Step profiling closes the loop. The included profiler fits quality curves from real measurements, so you can calibrate the model to your actual pipeline rather than guessing parameters.

Architecture

System flow diagram showing the optimization pipeline from step definition through quality modeling and water-fill allocation to the final output:

View architecture diagram on FigJam

graph LR
    A[Pipeline Definition] --> B[Step Quality Curves]
    B --> C[Water-Fill Optimizer]
    D[Total Token Budget] --> C
    C --> E[Marginal Gain Loop]
    E --> F[SLSQP Polish]
    F --> G[Optimal Allocation]
    G --> H[Sensitivity Report]

Loading

Each step's quality is modeled as q(t) = alpha * (1 - exp(-beta * t)), a saturating exponential where alpha is the quality ceiling and beta controls how fast quality rises with more tokens. End-to-end quality is the geometric mean of per-step effective qualities, with upstream error propagation penalizing downstream steps proportionally to upstream quality shortfall.

Results

The chart compares end-to-end pipeline quality for uniform allocation (gray) versus optimized allocation (blue) across seven total token budgets on a 4-step RAG pipeline (retrieval context stuffing, summarization, chain-of-thought generation, validation).

The metric is the geometric mean of effective per-step qualities, where each step's effective quality includes a propagation penalty from upstream steps. The baseline is uniform allocation (total budget divided equally among all steps). The optimized allocation is found by pipealloc's water-filling optimizer.

Key finding: at a 4096-token budget, uniform allocation gives each step 1024 tokens. The optimizer instead gives generation 2511 tokens (it has a slow-ramping quality curve and benefits most from extra budget), retrieval 570, summarization 800, and validation only 215 (it saturates at around 230 tokens). This reallocation improves end-to-end quality from 0.609 to 0.718, a 17.9% gain at zero additional token cost. Peak improvement across all tested budgets is 17.9% at 4096 tokens. Even at the tightest budget (512 tokens), the optimizer finds a 1.6% gain.

Budget	Uniform Quality	Optimized Quality	Gain
512	0.1923	0.1954	+1.6%
1024	0.3327	0.3568	+7.3%
2048	0.4825	0.5488	+13.8%
3072	0.5599	0.6538	+16.8%
4096	0.6088	0.7180	+17.9%
6144	0.6716	0.7866	+17.1%
8192	0.7127	0.8179	+14.8%

How to Run

# Clone and install
git clone https://github.com/joelvarun/pipealloc.git
cd pipealloc
pip install -r requirements.txt

# Run the demo (generates docs/before_after.png and prints results)
python examples/demo.py

Using pipealloc in your own pipeline

from pipealloc import PipelineStep, Pipeline, BudgetOptimizer

# Define your pipeline steps with quality curve parameters
pipeline = Pipeline(propagation_gamma=0.5)

pipeline.add_step(PipelineStep(
    name="retrieval",
    min_tokens=100, max_tokens=4000,
    alpha=0.90, beta=0.007,
    propagation_weight=1.3,
))
pipeline.add_step(PipelineStep(
    name="generation",
    min_tokens=100, max_tokens=8000,
    alpha=0.96, beta=0.0004,
    propagation_weight=0.7,
))

# Optimize for a given budget
optimizer = BudgetOptimizer(pipeline)
result = optimizer.optimize(total_budget=4096)

print(result.summary())
# result.allocations = {"retrieval": 620, "generation": 3476}
# result.improvement_pct = +15.2%

Calibrating from real measurements

If you have quality measurements from running your pipeline at different token budgets, use the profiler to fit quality curves automatically:

from pipealloc import StepProfiler

profiler = StepProfiler("generation", min_tokens=100, max_tokens=8000)

# Add measurements: (token_budget, measured_quality_score)
profiler.add_measurements([
    (200, 0.12),
    (500, 0.28),
    (1000, 0.41),
    (2000, 0.59),
    (4000, 0.78),
])

# Fit returns a PipelineStep with calibrated alpha and beta
step = profiler.fit(propagation_weight=0.7)
print(f"alpha={step.alpha:.3f}, beta={step.beta:.6f}")

# Check fit quality
residuals = profiler.residuals(step)
print(f"RMSE: {residuals['rmse']:.4f}")

Example Output

Full output from python examples/demo.py:

Pipeline with 4 steps (gamma=0.5):
  [1] retrieval_context: tokens=[100, 4000], alpha=0.90, beta=0.0070, saturation=658
  [2] summarization: tokens=[100, 3000], alpha=0.88, beta=0.0040, saturation=1152
  [3] generation: tokens=[100, 8000], alpha=0.96, beta=0.0004, saturation=8000
  [4] validation: tokens=[50, 1000], alpha=0.82, beta=0.0200, saturation=231

Optimization result (waterfill):
  Total budget: 4096 tokens
  Tokens used:  4096
  End-to-end quality (optimized): 0.7180
  End-to-end quality (uniform):   0.6088
  Improvement: +17.9%
  Per-step allocation and quality:
    retrieval_context: 570 tokens, quality=0.8834
    summarization: 800 tokens, quality=0.7934
    generation: 2511 tokens, quality=0.5490
    validation: 215 tokens, quality=0.6906

Sensitivity analysis (4096-token budget, 15% perturbation):
  Quality drop when each step is under-allocated:
    generation: -0.0102 (1.02% quality lost)
    summarization: -0.0024 (0.24% quality lost)
    retrieval_context: -0.0020 (0.20% quality lost)
    validation: -0.0007 (0.07% quality lost)

The sensitivity analysis confirms that generation is the bottleneck: reducing its allocation by 15% causes the largest quality drop, while reducing validation barely matters.