adaptive-k 0.1.1

Entropy-guided dynamic expert selection for MoE models. Reduce inference costs by 30-50%.

Adaptive-K SDK

Entropy-guided dynamic expert selection for Mixture-of-Experts models
Reduce inference costs by 30-50% with proven methodology.

---

🚀 Quick Start

pip install adaptive-k

from adaptive_k import AdaptiveKRouter

# Load pre-calibrated router
router = AdaptiveKRouter.from_pretrained("mixtral-8x7b")

# Route tokens
indices, weights, metrics = router.route(router_logits, return_metrics=True)

print(f"Compute savings: {metrics.compute_savings:.1%}")
# Output: Compute savings: 47.2%

---

📊 Proven Results

Model	Savings	Quality Retained
Mixtral 8x7B	52.5%	99.8%
Qwen-MoE	32.4%	99.9%
OLMoE-1B-7B	24.7%	99.7%

---

💡 How It Works

Adaptive-K dynamically selects the number of experts (K) based on routing entropy:

Low entropy (confident) → K=1 → 87.5% compute saved
Medium entropy         → K=2 → 75% compute saved  
High entropy (uncertain) → K=4 → Full routing

The key insight: when the router is confident, fewer experts are needed.

---

📖 Usage

Basic Routing

from adaptive_k import AdaptiveKRouter

router = AdaptiveKRouter.from_pretrained("mixtral-8x7b")

# Your MoE router logits (batch_size, num_experts)
router_logits = model.router(hidden_states)

# Adaptive-K routing
expert_indices, expert_weights, _ = router.route(router_logits)

# Use selected experts
output = execute_experts(hidden_states, expert_indices, expert_weights)

Custom Calibration

from adaptive_k import Calibrator

calibrator = Calibrator(
    target_savings=0.40,      # 40% target savings
    quality_threshold=0.99    # Max 1% quality loss
)

result = calibrator.calibrate(
    model=your_model,
    dataset=calibration_data
)

print(f"Optimal thresholds: {result.optimal_thresholds}")
print(f"Expected savings: {result.expected_savings:.1%}")

Check Statistics

# After processing many tokens
print(router.stats)
# {
#   'tokens_processed': 1_234_567,
#   'average_savings': 0.472,
#   'estimated_cost_reduction': '47.2%'
# }

---

🔧 Configuration

from adaptive_k import AdaptiveKRouter, RoutingConfig

config = RoutingConfig(
    k_values=[1, 2, 4],           # Available K values
    entropy_thresholds=[0.6, 1.2], # H < 0.6 → K=1, H < 1.2 → K=2, else K=4
    num_experts=8
)

router = AdaptiveKRouter(config=config)

---

🔌 Integrations

HuggingFace Transformers

# Coming in v0.2.0
router = AdaptiveKRouter.from_pretrained("mixtral-8x7b")
model = router.patch(model)  # Automatic integration

vLLM

# Coming in v0.3.0
from adaptive_k.integrations import vllm_patch
model = vllm_patch(model, router)

TensorRT-LLM

See our TensorRT-LLM PR #10672 for native integration.

---

📈 Benchmarking

# CLI benchmark
adaptive-k benchmark --model mixtral-8x7b --dataset wikitext-2

# Output:
# Model: mixtral-8x7b
# Dataset: wikitext-2
# Baseline perplexity: 5.42
# Adaptive-K perplexity: 5.44 (+0.4%)
# Compute savings: 47.2%

---

📄 License

Apache 2.0 - Free for commercial use.

---

🔗 Links

• Website: https://adaptive-k.vertexdata.it
• Paper: Entropy-Guided Dynamic Expert Selection
• GitHub: https://github.com/Gabrobals/sbm-efficient

---

📞 Support

• Email: [e-mail]
• Issues: GitHub Issues

---

Made with ❤️ by Vertex Data