SGLang serving: external model package, launcher, client, benchmark; card section

README.md

@@ -64,7 +64,7 @@ a small MLP head on the last-token latent (`mlp_heads_35b/`, one head per task,
 * `modeling_qwen35_moe_seqcls.py` — `Qwen3_5MoeForSequenceClassification` for the 35B-A3B backbone.
 * `mlp_heads_35b/<task>/` — `head.pt` + `norm.npz` + `meta.json`, the 35B latent + MLP heads behind the second radar.
 * `code/` — everything used here: the trainer and data mixture builder, the evaluation harness, Flappy Bird, Doom
-  (text and pixels), the Minecraft scaffold and bot, the radar.
+  (text and pixels), the Minecraft scaffold and bot, the radar, the SGLang package / launcher / client / benchmark.
 * `videos/` — Flappy Bird, Doom and Minecraft replays; `results/` — raw JSON for every run and the full report.
 
 ## Use it
@@ -104,4 +104,41 @@ text = model.config.nli_template.format(premise="...", hypothesis="...")
 Images go inside the premise as `<|vision_start|><|image_pad|>…<|vision_end|>` with `pixel_values` / `image_grid_thw`
 from the Qwen3.5 image processor; see `code/doom_vision.py` and `code/eval_image_nli.py`.
 
+## Serve it with SGLang
+
+SGLang has no sequence-classification class for Qwen3.5, so `code/sglang_openjev/` is an external model package
+that adds the `score` head to SGLang's `Qwen3_5ForConditionalGeneration` (hybrid cache, mrope and the vision tower
+stay as they are). Text and images both work; `/classify` returns the three raw logits.
+
+```bash
+hf download AlexWortega/openjev --include "qwen3.5-0.8b-nli-v2s-long/*" "code/*" --local-dir openjev
+cd openjev/code && bash serve_sglang.sh ../qwen3.5-0.8b-nli-v2s-long 30000      # tested with sglang 0.5.19
+```
+
+```python
+from sglang_client import OpenJevSGLang          # code/sglang_client.py
+jev = OpenJevSGLang("http://127.0.0.1:30000")
+jev.predict([("A man is playing a guitar.", "Someone is making music.")])    # [[con, ent, neu]]
+jev.predict([("A photograph of a scene:", "There is a dog.")], images=["dog.jpg"])
+```
+
+Without the script: `SGLANG_EXTERNAL_MODEL_PACKAGE=sglang_openjev` with `code/` on `PYTHONPATH`, then
+`python -m sglang.launch_server --model-path <dir> --is-embedding --json-model-override-args
+'{"architectures": ["Qwen3_5ForConditionalGeneration"]}'`.
+
+Same predictions as transformers, 1.5-3x the throughput. `qwen3.5-0.8b-nli-v2s-long` on one RTX A6000 that was
+shared with another job (so absolute speed is a lower bound), transformers at batch 32 (`code/bench_sglang.py`):
+
+| task | pairs | tokens / pair | acc SGLang | acc transformers | SGLang pairs/s | transformers pairs/s | speed-up |
+|---|---|---|---|---|---|---|---|
+| MNLI m+mm | 19647 | 44 | 86.63 | 86.64 | 321 | 152 | 2.1x |
+| ANLI r1 | 1000 | 105 | 65.3 | 64.9 | 189 | 119 | 1.6x |
+| ANLI r2 | 1000 | 103 | 50.6 | 50.8 | 181 | 110 | 1.6x |
+| ANLI r3 | 1200 | 92 | 48.6 | 48.6 | 218 | 95 | 2.3x |
+| WANLI | 5000 | 40 | 73.7 | 73.7 | 400 | 196 | 2.0x |
+| SciTail | 2126 | 43 | 93.3 | 93.4 | 319 | 217 | 1.5x |
+| ConTRoL (long) | 805 | 612 | 51.9 | 51.4 | 45 | 15 | 3.1x |
+| ARC-Challenge (rerank) | 9374 | 49 | 49.2 | 49.2 | 334 | 197 | 1.7x |
+| HellaSwag (rerank, 2k questions) | 16000 | 114 | 37.6 | 37.5 | 176 | 109 | 1.6x |
+
 Reference point: [dleemiller's NLI cross-encoders](https://huggingface.co/blog/dleemiller/nli-xenc-ways-to-use). Licence MIT.

code/bench_sglang.py

@@ -0,0 +1,95 @@
+"""Accuracy + throughput of an openjev checkpoint: transformers (eval.py's NLIScorer) vs the SGLang server.
+
+    ./serve_sglang.sh ckpt/qwen3.5-0.8b-nli-v2s-long &
+    python bench_sglang.py --ckpt ckpt/qwen3.5-0.8b-nli-v2s-long --out results/sglang/bench.json
+
+Both backends go through the same eval.py task code; only `predict` is timed (no dataset loading)."""
+import argparse
+import json
+import os
+import time
+from types import SimpleNamespace
+
+import numpy as np
+import torch
+
+import eval as E
+from sglang_client import OpenJevSGLang
+
+
+class Timed:
+    def __init__(self, predict, tok, max_len):
+        self._predict, self.tok, self.max_len = predict, tok, max_len
+        self.reset()
+
+    def reset(self):
+        self.sec, self.pairs, self.tokens = 0.0, 0, 0
+
+    def predict(self, pairs):
+        t = time.perf_counter()
+        out = self._predict(pairs)
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        self.sec += time.perf_counter() - t
+        self.pairs += len(pairs)
+        texts = [E_TEMPLATE.format(premise=p.strip(), hypothesis=h.strip()) for p, h in pairs]
+        self.tokens += sum(min(len(x), self.max_len) for x in self.tok(texts)["input_ids"])
+        return out
+
+
+def run(scorer, tasks, mc_items, mnli_n):
+    res = {}
+    for t in tasks:
+        scorer.reset()
+        if t == "mnli":
+            r = E.eval_mnli(scorer, mnli_n)
+            acc = float(np.mean([v["acc"] for v in r.values()]))
+        elif t in E.NLI_SETS:
+            acc = E.eval_nli_set(scorer, t)["acc"]
+        else:
+            acc = E.eval_mc(scorer, mc_items[t])["rerank_acc"]
+        res[t] = {"acc": acc, "pairs": scorer.pairs, "tokens": scorer.tokens, "sec": round(scorer.sec, 2),
+                  "pairs_per_s": round(scorer.pairs / scorer.sec, 1), "tok_per_s": round(scorer.tokens / scorer.sec)}
+        print(t, res[t], flush=True)
+    return res
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--url", default="http://127.0.0.1:30000")
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--tasks", nargs="+", default=["mnli", "anli_r1", "anli_r2", "anli_r3", "wanli", "scitail", "control",
+                                                   "arc_challenge", "hellaswag"])
+    ap.add_argument("--backends", nargs="+", default=["sglang", "hf"])
+    ap.add_argument("--mnli-n", type=int, default=None)
+    ap.add_argument("--mc-n", type=int, default=2000, help="hellaswag / mmlu subsample")
+    ap.add_argument("--bs", type=int, default=32)
+    ap.add_argument("--max-len", type=int, default=4096)
+    ap.add_argument("--workers", type=int, default=16)
+    a = ap.parse_args()
+    from transformers import AutoConfig, AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(a.ckpt)
+    E_TEMPLATE = AutoConfig.from_pretrained(a.ckpt).nli_template
+    margs = SimpleNamespace(mc_n=a.mc_n, chess_n=500, fewshot=5)
+    mc_items = {t: E.MC_TASKS[t](margs) for t in a.tasks if t in E.MC_TASKS}
+    results = {}
+    for b in a.backends:
+        print(f"\n===== {b}")
+        if b == "hf":
+            model = E.NLIScorer(a.ckpt, bs=a.bs, max_len=a.max_len)
+            predict = model.predict
+        else:
+            predict = OpenJevSGLang(a.url, template=E_TEMPLATE, bs=a.bs, workers=a.workers).predict
+        predict([("warm", "up")] * a.bs)
+        results[b] = run(Timed(predict, tok, a.max_len), a.tasks, mc_items, a.mnli_n)
+    os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True)
+    json.dump(results, open(a.out, "w"), indent=2)
+    print("\n| task | pairs | avg tok | " + " | ".join(f"{b} acc | {b} pairs/s" for b in a.backends) + " | speedup |")
+    print("|---|---|---|" + "---|---|" * len(a.backends) + "---|")
+    for t in a.tasks:
+        r0 = results[a.backends[0]][t]
+        row = " | ".join(f"{results[b][t]['acc']:.4f} | {results[b][t]['pairs_per_s']}" for b in a.backends)
+        sp = r0["pairs_per_s"] / results[a.backends[-1]][t]["pairs_per_s"] if len(a.backends) > 1 else 1.0
+        print(f"| {t} | {r0['pairs']} | {r0['tokens'] // r0['pairs']} | {row} | {sp:.1f}x |")

code/serve_sglang.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+# Serve an openjev checkpoint with SGLang. Usage: ./serve_sglang.sh [ckpt dir or HF snapshot dir] [port] [extra sglang args]
+# Needs a local dir with config.json at its root (for the Hub repo: hf download AlexWortega/openjev --include
+# "qwen3.5-0.8b-nli-v2s-long/*" --local-dir openjev_hf). Query with sglang_client.py.
+# ATTN_BACKEND defaults to triton: flashinfer JIT-compiles with the system nvcc and fails on nvcc < 12.4.
+# MEM_FRACTION is a share of the WHOLE card, other processes included: on a shared GPU set it close to
+# 1 - (wanted headroom / total), e.g. 0.85 on a 48GB card that already has 32GB taken.
+HERE=$(cd "$(dirname "$0")" && pwd)
+MODEL=${1:-ckpt/qwen3.5-0.8b-nli-v2s-long}
+PORT=${2:-30000}
+PY=${PY:-$HOME/venvs/sglang/bin/python}
+export PATH=$(dirname "$PY"):$PATH  # ninja for sglang's JIT kernels
+export PYTHONPATH=$HERE:$PYTHONPATH
+export SGLANG_EXTERNAL_MODEL_PACKAGE=sglang_openjev
+exec $PY -m sglang.launch_server --model-path "$MODEL" --port "$PORT" --host 0.0.0.0 \
+  --is-embedding --json-model-override-args '{"architectures": ["Qwen3_5ForConditionalGeneration"]}' \
+  --attention-backend ${ATTN_BACKEND:-triton} --served-model-name openjev ${MEM_FRACTION:+--mem-fraction-static $MEM_FRACTION} "${@:3}"

code/sglang_client.py

@@ -0,0 +1,92 @@
+"""Client for an openjev checkpoint served by SGLang (serve_sglang.sh). Same interface as
+`OpenJevCrossEncoder.predict/rerank`, plus `--check` to compare the server against transformers.
+
+    python sglang_client.py --premise "A man is playing a guitar." --hypothesis "Someone is making music."
+    python sglang_client.py --check ckpt/qwen3.5-0.8b-nli-v2s-long        # max |p_sglang - p_hf| on a few pairs
+"""
+from __future__ import annotations
+
+import argparse
+import base64
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+import requests
+
+CON, ENT, NEU = 0, 1, 2
+LABELS = ["contradiction", "entailment", "neutral"]
+DEFAULT_TEMPLATE = "Premise: {premise}\nHypothesis: {hypothesis}"
+IMG_MARK = "<<IMG>>"
+IMG_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>"  # SGLang expands the single pad to the patch count
+
+
+class OpenJevSGLang:
+    def __init__(self, url: str = "http://127.0.0.1:30000", template: str = DEFAULT_TEMPLATE, bs: int = 32,
+                 workers: int = 16):
+        # `workers` requests of `bs` pairs are in flight at once; the server batches across them
+        self.url, self.template, self.bs, self.workers = url.rstrip("/"), template, bs, workers
+
+    def _post(self, payload):
+        r = requests.post(f"{self.url}/classify", json=payload, timeout=600)
+        r.raise_for_status()
+        out = r.json()
+        return [o["embedding"] for o in (out if isinstance(out, list) else [out])]
+
+    def logits(self, pairs, images=None) -> np.ndarray:
+        """Raw [contradiction, entailment, neutral] logits. `images[i]` is a path / URL / bytes or None and is
+        spliced into premise i at the `<<IMG>>` marker (appended if the marker is absent)."""
+        payloads = []
+        for i in range(0, len(pairs), self.bs):
+            texts, imgs = [], []
+            for j, (p, h) in enumerate(pairs[i:i + self.bs]):
+                img = images[i + j] if images else None
+                if img is not None:
+                    p = p.replace(IMG_MARK, IMG_BLOCK) if IMG_MARK in p else f"{p.rstrip()} {IMG_BLOCK}"
+                    img = base64.b64encode(img).decode() if isinstance(img, bytes) else img
+                texts.append(self.template.format(premise=p.strip(), hypothesis=h.strip()))
+                imgs.append(img)
+            payload = {"text": texts}
+            if any(x is not None for x in imgs):
+                payload["image_data"] = imgs
+            payloads.append(payload)
+        with ThreadPoolExecutor(self.workers) as ex:
+            out = [z for part in ex.map(self._post, payloads) for z in part]
+        return np.asarray(out, np.float32)
+
+    def predict(self, pairs, images=None) -> np.ndarray:
+        """Softmax probabilities [contradiction, entailment, neutral] per pair."""
+        z = self.logits(pairs, images)
+        z = np.exp(z - z.max(-1, keepdims=True))
+        return z / z.sum(-1, keepdims=True)
+
+    def rerank(self, question: str, options, hyp_fmt: str = "The correct answer is: {}") -> int:
+        return int(self.predict([(question, hyp_fmt.format(o)) for o in options])[:, ENT].argmax())
+
+
+CHECK_PAIRS = [
+    ("A man is playing a guitar.", "Someone is making music."),
+    ("A man is playing a guitar.", "The man is asleep."),
+    ("A man is playing a guitar.", "The man is a professional musician."),
+    ("The Eiffel Tower is in Paris. " * 40, "The Eiffel Tower is in Berlin."),
+    ("What is the capital of France?", "The correct answer is: Paris"),
+]
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--url", default="http://127.0.0.1:30000")
+    ap.add_argument("--premise")
+    ap.add_argument("--hypothesis")
+    ap.add_argument("--image")
+    ap.add_argument("--check", metavar="CKPT", help="compare against transformers on this checkpoint")
+    a = ap.parse_args()
+    client = OpenJevSGLang(a.url)
+    if a.check:
+        from modeling_openjev import OpenJevCrossEncoder
+
+        ours, ref = client.predict(CHECK_PAIRS), OpenJevCrossEncoder(a.check).predict(CHECK_PAIRS)
+        for (p, h), x, y in zip(CHECK_PAIRS, ours, ref):
+            print(f"{h[:50]:52s} sglang {np.round(x, 4)}  hf {np.round(y, 4)}")
+        print("max abs diff:", float(np.abs(ours - ref).max()), "| argmax agree:", bool((ours.argmax(1) == ref.argmax(1)).all()))
+    else:
+        p = client.predict([(a.premise, a.hypothesis)], [a.image] if a.image else None)[0]
+        print({l: round(float(v), 4) for l, v in zip(LABELS, p)})

code/sglang_openjev/__init__.py

@@ -0,0 +1,2 @@
+"""SGLang external model package for openjev. Enable with SGLANG_EXTERNAL_MODEL_PACKAGE=sglang_openjev
+(this directory's parent must be on PYTHONPATH); see serve_sglang.sh."""

code/sglang_openjev/qwen3_5_seqcls.py

@@ -0,0 +1,48 @@
+"""Qwen3_5ForSequenceClassification for SGLang (which ships no seq-cls head for Qwen3.5).
+
+SGLang keys the hybrid GDN cache, the mrope setup and the Qwen-VL multimodal processor on the architecture
+*name* `Qwen3_5ForConditionalGeneration`, so instead of registering a new name this module overrides that
+entry (the registry loads external packages with overwrite=True) and the server is launched with
+`--json-model-override-args '{"architectures": ["Qwen3_5ForConditionalGeneration"]}' --is-embedding`.
+Generation is untouched; with get_embedding the last-token hidden state goes through the linear `score`
+head, same pooling as transformers (last non-pad token), and the raw logits come back as the "embedding".
+"""
+from typing import Optional
+
+import torch
+from torch import nn
+
+from sglang.srt.layers.pooler import Pooler, PoolingType, score_and_pool
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.models import qwen3_5 as _qwen3_5
+from sglang.srt.models.qwen3_vl import general_mm_embed_routine
+
+
+class Qwen3_5ForConditionalGeneration(_qwen3_5.Qwen3_5ForConditionalGeneration):
+    def __init__(self, config, quant_config: Optional[QuantizationConfig] = None, prefix: str = ""):
+        super().__init__(config, quant_config, prefix)
+        # self.config is the text config from here on; num_labels lives on the top-level one
+        self.score = nn.Linear(self.config.hidden_size, config.num_labels, bias=False)
+        self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=False)
+
+    @torch.no_grad()
+    def forward(self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch,
+                get_embedding: bool = False, pp_proxy_tensors=None):
+        if not get_embedding:
+            return super().forward(input_ids, positions, forward_batch, False, pp_proxy_tensors)
+        if self.is_mrope_enabled:
+            positions = forward_batch.mrope_positions
+        hidden_states = general_mm_embed_routine(
+            input_ids=input_ids,
+            forward_batch=forward_batch,
+            language_model=self.model,
+            multimodal_model=self,
+            positions=positions,
+            use_deepstack=self.use_deepstack,
+            pp_proxy_tensors=pp_proxy_tensors,
+        )
+        return score_and_pool(self.score, self.pooler, hidden_states, forward_batch, input_ids)
+
+
+EntryClass = [Qwen3_5ForConditionalGeneration]