tensorbend-predictor

Cross-layer expert predictor for Mixture-of-Experts models. Predicts which experts the router will select before the gate runs, enabling prefetch of expert weights from NVMe/DRAM or pre-staging of EP all2all buffers.

Trains a lightweight MLP per MoE layer that takes the pre-MoE hidden states + previous layer's routing decision as input, achieving 99.8% training accuracy and 96.9% live prefill accuracy on DeepSeek-style MoE architectures.

Step-by-step reproduction

These instructions reproduce the validated results on 0xSero/Kimi-K2.5-PRISM-REAP-72 with 8x H100.

1. Setup

git clone https://github.com/eelbaz/tensorbend-predictor.git
cd tensorbend-predictor
pip install -r requirements.txt

Verify vLLM can load the model:

vllm serve 0xSero/Kimi-K2.5-PRISM-REAP-72 \
    --tensor-parallel-size 8 \
    --max-model-len 2048 \
    --trust-remote-code \
    --port 8000 &

# Wait for "Application startup complete", then test:
curl http://localhost:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{"model":"0xSero/Kimi-K2.5-PRISM-REAP-72","messages":[{"role":"user","content":"Say hello"}],"max_tokens":10}'

# Kill the server after testing:
pkill -f "vllm serve"

2. Download the model (if not using HF hub directly)

huggingface-cli download 0xSero/Kimi-K2.5-PRISM-REAP-72 \
    --local-dir /path/to/kimi-reap-72

3. Train the predictor

VLLM_ALLOW_INSECURE_SERIALIZATION=1 \
python scripts/collect_and_train.py \
    --model-dir /path/to/kimi-reap-72 \
    --output predictor_reap72.pt \
    --n-gpus 8 \
    --predictor-dim 256

This takes ~25 minutes:

• Phase 1 (~3 min): loads model via vLLM offline inference
• Phase 2 (~1 sec): installs gate forward hooks
• Phase 3 (~20 min): runs 100 calibration prompts, collects routing data
• Phase 4 (~2 min): trains 60 per-layer predictors

Expected output:

Accuracy: avg=99.8% min=99.4% max=100.0%
Saved: predictor_reap72.pt (450MB)

4. Verify with live vLLM inference

The repo includes a pre-trained predictor (predictor_reap72_dim256.pt) you can use directly:

import torch
from predictor import ExpertPredictorSet

# Load predictor
pset = ExpertPredictorSet(
    num_layers=60, hidden_dim=7168, num_experts=72,
    predictor_dim=256, device="cuda:0", use_prev_routing=True,
)
pset.load("predictor_reap72_dim256.pt")

# Predict experts for a given hidden state
h = torch.randn(1, 7168)  # pre-MoE hidden state
predicted_experts = pset.predict(layer_idx=0, hidden_states=h, top_k=8)
print(f"Predicted experts: {predicted_experts}")

Adapting to other models

The predictor works with any DeepSeek-style MoE model served by vLLM. Key parameters to adjust:

These are auto-detected from the model config. No manual changes needed.

# REAP-530B (192 experts)
python scripts/collect_and_train.py \
    --model-dir Ex0bit/Kimi-K2.5-PRISM-REAP-530B-A32B \
    --output predictor_reap530b.pt \
    --n-gpus 8

GPU configurations

Choosing predictor_dim

--predictor-dim 256 is the recommended default.

Validated results

Trained and live-validated on 0xSero/Kimi-K2.5-PRISM-REAP-72 (8x H100 80GB):

Key findings:

• Pure BCE outperforms ranking loss + importance weighting (99.8% vs 86%)
• 100 diverse prompts > 30 prompts with longer generation (generalization matters)
• Prefill accuracy tracks training accuracy; decode is an architectural ceiling for feedforward predictors

Architecture

predictor/
    __init__.py          # Public API
    model.py             # ExpertPredictor, ExpertPredictorSet
    train.py             # Training (pure BCE, cosine LR)
    hooks.py             # vLLM gate hook for data collection
scripts/
    collect_and_train.py # End-to-end training
paper.pdf                # Technical report
predictor_reap72_dim256.pt  # Pre-trained weights (Git LFS)

How the hooks work

A single gate forward hook captures both:

1. Gate input = exact hidden states the gate kernel receives
2. Gate output = exact router logits from the kernel

This guarantees the predictor trains on the same data the gate uses, with zero numerical mismatch. The hooks use vLLM's apply_model() API and only collect on TP rank 0.

Paper

See paper.pdf for the full technical report including:

• Architecture details and design decisions
• Training methodology and loss function ablation
• End-to-end performance on DGX Spark (1.77 tok/s, 9.5 GB GPU)
• Comparison to 15+ related works (2024-2026)

License

Apache-2.0