A small output-length ranker used to schedule LLM inference requests shortest-predicted-first (SJF), the mechanism behind the Learning-to-Rank (LTR) scheduler of Fu et al. (NeurIPS 2024) and the FDU latency study (Kumar et al.). Given a chat prompt it emits one scalar score — higher score ⇒ shorter expected output ⇒ run this request earlier — which is mapped to a vLLM request priority. It is the B1 ranker in the capstone KVCache-Coordinated Latency Optimization.

Architecture: facebook/opt-125m backbone (AutoModel, fp32) + masked-mean pooling over the prompt tokens + a Linear(768 → 1) scoring head.
Objective: ListMLE listwise ranking loss (list size 16), seed 0, 10 epochs.
Labels: real output lengths from LMSYS-Chat-1M (the FDU paper's methodology).
Target serve model: Llama-3.1-8B-Instruct (the ranker itself is model-agnostic — it ranks prompts, not tokens).

Training


Base	`facebook/opt-125m`
Loss	ListMLE (listwise), list size 16
Epochs	10 (resumed 5→10, continuous curve), seed 0
Data	LMSYS-Chat-1M prompt → output-length lists
ListMLE loss	25.67 → 23.04 → 19.98 → 18.02 → 16.22 → 15.01 → 13.84 → 12.92 → 12.51 → 11.67
Ranking quality	Kendall τ ≈ 0.71 on a held-out LMSYS split (score vs true length)

Monotone loss with no overfit rebound at 10 epochs (the paper's balance point).

The target-sampled output-length labels are published separately: nvmmonkey/llama31-8b-output-lengths (LMSYS prompts withheld per license).

Usage

The scoring head is custom, so load the state dict onto the architecture shipped here (modeling_opt_ranker.py):

python

import torch, importlib.util
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from transformers import AutoTokenizer

repo = snapshot_download("<your-hf-user>/opt125m-ltr-ranker")

spec = importlib.util.spec_from_file_location("m", f"{repo}/modeling_opt_ranker.py")
m = importlib.util.module_from_spec(spec); spec.loader.exec_module(m)

ranker = m.build_ranker("facebook/opt-125m")          # OPT backbone + linear score head
ranker.load_state_dict(load_file(f"{repo}/model.safetensors"))
ranker.eval()

tok = AutoTokenizer.from_pretrained(repo)
enc = tok(["Explain quantum computing in one sentence.",
           "Write a 2000-word essay on the French Revolution."],
          return_tensors="pt", padding=True, truncation=True, max_length=512)

with torch.no_grad():
    scores = ranker(enc.input_ids, enc.attention_mask)  # higher => shorter output => schedule first
print(scores)   # the short prompt should score higher than the long one

To serve with it, stamp each request's vLLM priority from the score (see the code repo's ltr/scheduler/priority.py and serving/serve_b1_ltr.sh, launched with --scheduling-policy priority).

Provenance, license & citation

Base model facebook/opt-125m is under the OPT license (research / non-commercial); this fine-tune inherits it.
Training labels come from LMSYS-Chat-1M, under the LMSYS-Chat-1M dataset license — non-commercial, agree to its terms before use.
Method reproduces prior work — please cite it:
- Y. Fu et al., "Efficient LLM scheduling by learning to rank," NeurIPS 2024. — github.com/hao-ai-lab/vllm-ltr
- A. Saravana Kumar et al., "An empirical study on latency reduction techniques for large language models," FDU, 2026.

Trained for the FDU CSCI 6806 capstone (Guoliang Liu, Wenhui Kang, Junpeng Huang).

nvmmonkey/opt125m-ltr-ranker

OPT-125M Output-Length Ranker (Learning-to-Rank scheduling for LLM serving)

Training

Usage

Provenance, license & citation

nvmmonkey/
opt125m-ltr-ranker