sgl-router: experimental Rust HTTP router for SGLang worker pools (#25851)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
aae04b1241
commit
6e8fe176be
@@ -0,0 +1,237 @@
|
||||
"""
|
||||
Generator + validator for KV-event block-hash parity fixtures.
|
||||
|
||||
Two modes:
|
||||
|
||||
python3 experimental/sgl-router/tests/scripts/generate_kv_events_hash_parity.py
|
||||
Regenerate the committed JSON fixture from the locally-replicated
|
||||
algorithm. Run this when changing block-hash logic or adding new
|
||||
shape coverage. CI's drift-check step runs this in --check mode.
|
||||
|
||||
python3 experimental/sgl-router/tests/scripts/generate_kv_events_hash_parity.py --validate-against-sglang
|
||||
Import the real `sglang.srt.mem_cache.radix_cache.RadixKey.hash_page`
|
||||
and assert it agrees with the locally-replicated algorithm on every
|
||||
fixture case. This is the only place the replica and the real
|
||||
SGLang implementation are checked against each other. Run it
|
||||
nightly (or whenever sglang is available on the Python path).
|
||||
|
||||
# Authority
|
||||
|
||||
Source-of-truth implementation:
|
||||
- `python/sglang/srt/mem_cache/radix_cache.py::RadixKey.hash_page`
|
||||
- `python/sglang/srt/mem_cache/utils.py::hash_str_to_int64`
|
||||
|
||||
`hash_page_chain` below replicates that algorithm verbatim (no `import
|
||||
sglang`) so the script runs without the heavy SGLang dependency tree and
|
||||
can be audited at a glance. The algorithm is intentionally tiny:
|
||||
|
||||
sha256(prior_digest_bytes ++ token_LE_u32 ++ token_LE_u32 ++ ...)
|
||||
truncate to i64 = signed(first 16 hex chars)
|
||||
|
||||
If SGLang ever changes the algorithm, update both the SGLang side AND
|
||||
this script in the same commit; the Rust port in
|
||||
`src/policies/kv_events/hash.rs` will then need the corresponding
|
||||
update. The nightly `--validate-against-sglang` job is the safety net
|
||||
that catches an SGLang-side change the human forgot to mirror here.
|
||||
|
||||
# Output format
|
||||
|
||||
A JSON array of cases. Each case is:
|
||||
{
|
||||
"name": "<descriptive label>",
|
||||
"tokens": [<u32>, ...],
|
||||
"block_size": <usize>,
|
||||
"expected_i64_hashes": [<i64>, ...]
|
||||
}
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import pathlib
|
||||
import sys
|
||||
|
||||
|
||||
def hash_page_chain(tokens: list[int], block_size: int) -> list[int]:
|
||||
"""Compute the i64-truncated block hashes for `tokens` using SGLang's
|
||||
`RadixKey.hash_page` algorithm + `hash_str_to_int64`.
|
||||
|
||||
Returns one i64 per full or partial block. A partial last block (when
|
||||
`len(tokens) % block_size != 0`) chains against the previous block's
|
||||
full 32-byte SHA256 digest, matching SGLang's behaviour.
|
||||
"""
|
||||
if block_size == 0:
|
||||
raise ValueError("block_size must be positive")
|
||||
|
||||
out: list[int] = []
|
||||
prior_digest: bytes | None = None
|
||||
n = len(tokens)
|
||||
if n == 0:
|
||||
return out
|
||||
# Walk every page boundary, including a trailing partial page.
|
||||
start = 0
|
||||
while start < n:
|
||||
end = min(start + block_size, n)
|
||||
hasher = hashlib.sha256()
|
||||
if prior_digest is not None:
|
||||
hasher.update(prior_digest)
|
||||
for t in tokens[start:end]:
|
||||
hasher.update(t.to_bytes(4, byteorder="little", signed=False))
|
||||
digest = hasher.digest()
|
||||
prior_digest = digest
|
||||
# hash_str_to_int64: first 16 hex chars (top 64 bits) -> signed i64.
|
||||
hex_digest = digest.hex()
|
||||
uint64_val = int(hex_digest[:16], 16)
|
||||
if uint64_val >= 2**63:
|
||||
i64 = uint64_val - 2**64
|
||||
else:
|
||||
i64 = uint64_val
|
||||
out.append(i64)
|
||||
start = end
|
||||
return out
|
||||
|
||||
|
||||
# Cases mirror the three existing `cross_language_golden_*` tests plus
|
||||
# additional shape coverage that exercises (a) zero-token edge, (b)
|
||||
# block_size = 1, (c) very long sequences, (d) odd boundaries.
|
||||
CASES: list[dict] = [
|
||||
{
|
||||
"name": "single_full_block",
|
||||
"tokens": [1, 2, 3, 4],
|
||||
"block_size": 4,
|
||||
},
|
||||
{
|
||||
"name": "partial_last_block",
|
||||
"tokens": [1, 2, 3, 4, 5],
|
||||
"block_size": 4,
|
||||
},
|
||||
{
|
||||
"name": "multi_block",
|
||||
"tokens": [10, 20, 30, 40, 50, 60, 70, 80],
|
||||
"block_size": 2,
|
||||
},
|
||||
{
|
||||
"name": "empty_tokens",
|
||||
"tokens": [],
|
||||
"block_size": 4,
|
||||
},
|
||||
{
|
||||
"name": "block_size_one",
|
||||
"tokens": [7, 8, 9],
|
||||
"block_size": 1,
|
||||
},
|
||||
{
|
||||
"name": "odd_boundary",
|
||||
"tokens": [100, 200, 300, 400, 500, 600, 700],
|
||||
"block_size": 3,
|
||||
},
|
||||
{
|
||||
"name": "long_sequence",
|
||||
# 128 tokens at block_size 16 → 8 blocks exactly.
|
||||
"tokens": list(range(1, 129)),
|
||||
"block_size": 16,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _materialize_cases() -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"name": c["name"],
|
||||
"tokens": c["tokens"],
|
||||
"block_size": c["block_size"],
|
||||
"expected_i64_hashes": hash_page_chain(c["tokens"], c["block_size"]),
|
||||
}
|
||||
for c in CASES
|
||||
]
|
||||
|
||||
|
||||
def _validate_against_sglang() -> int:
|
||||
"""Import the real SGLang `RadixKey.hash_page` and compare its output
|
||||
case-by-case against the locally-replicated `hash_page_chain`. Exits
|
||||
non-zero (and prints a diff-friendly summary) on any mismatch.
|
||||
|
||||
Returns 0 on success. This is the parity safety net for nightly CI.
|
||||
"""
|
||||
try:
|
||||
from sglang.srt.mem_cache.radix_cache import RadixKey
|
||||
except ImportError as e:
|
||||
print(
|
||||
f"--validate-against-sglang: cannot import sglang ({e}). "
|
||||
"Install sglang into the Python path before running this mode.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 2
|
||||
|
||||
failures: list[str] = []
|
||||
for c in CASES:
|
||||
local = hash_page_chain(c["tokens"], c["block_size"])
|
||||
if c["block_size"] == 0 or not c["tokens"]:
|
||||
# `RadixKey.hash_page` requires a non-empty page; the local
|
||||
# replica handles edge cases (empty input → empty list)
|
||||
# which the SGLang oracle would refuse. Skip these cases
|
||||
# under validation — the replica owns the boundary semantics.
|
||||
continue
|
||||
sglang_hashes: list[int] = []
|
||||
prior_hex: str | None = None
|
||||
for start in range(0, len(c["tokens"]), c["block_size"]):
|
||||
page = c["tokens"][start : start + c["block_size"]]
|
||||
key = RadixKey(token_ids=page, extra_key=None)
|
||||
hex_digest = key.hash_page(prior_hex)
|
||||
# SGLang's hash_page returns the hex digest; truncate to i64
|
||||
# the same way `hash_str_to_int64` does.
|
||||
uint64_val = int(hex_digest[:16], 16)
|
||||
i64 = uint64_val - (1 << 64) if uint64_val >= (1 << 63) else uint64_val
|
||||
sglang_hashes.append(i64)
|
||||
prior_hex = hex_digest
|
||||
if sglang_hashes != local:
|
||||
failures.append(f"case {c['name']}: local={local} sglang={sglang_hashes}")
|
||||
|
||||
if failures:
|
||||
print(
|
||||
"--validate-against-sglang: replica/SGLang DRIFT detected:",
|
||||
file=sys.stderr,
|
||||
)
|
||||
for f in failures:
|
||||
print(f" {f}", file=sys.stderr)
|
||||
return 1
|
||||
print(f"--validate-against-sglang: OK ({len(CASES)} cases agreed)")
|
||||
return 0
|
||||
|
||||
|
||||
def _write_fixture(cases_out: list[dict]) -> pathlib.Path:
|
||||
out_path = (
|
||||
pathlib.Path(__file__).resolve().parent.parent
|
||||
/ "fixtures"
|
||||
/ "kv_events_hash_parity.json"
|
||||
)
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with out_path.open("w") as f:
|
||||
json.dump(cases_out, f, indent=2, sort_keys=False)
|
||||
f.write("\n")
|
||||
return out_path
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--validate-against-sglang",
|
||||
action="store_true",
|
||||
help="Compare the local replica to the imported SGLang implementation "
|
||||
"and exit non-zero on drift. Requires sglang on the Python path.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.validate_against_sglang:
|
||||
return _validate_against_sglang()
|
||||
|
||||
cases_out = _materialize_cases()
|
||||
out_path = _write_fixture(cases_out)
|
||||
print(f"wrote {len(cases_out)} cases to {out_path}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,115 @@
|
||||
"""
|
||||
One-shot generator for tokenizer parity fixtures.
|
||||
|
||||
Run manually when adding a model or changing a prompt shape:
|
||||
python3 -m venv /tmp/parity-fixture-venv
|
||||
/tmp/parity-fixture-venv/bin/pip install transformers
|
||||
/tmp/parity-fixture-venv/bin/python experimental/sgl-router/tests/scripts/generate_parity_fixtures.py
|
||||
|
||||
CI does NOT run this — it consumes the committed JSON.
|
||||
|
||||
Model substitutions (gated models → public siblings of same family):
|
||||
- Qwen/Qwen3-30B-A3B (gated) → Qwen/Qwen3-0.6B (same Qwen3 family, public)
|
||||
- deepseek-ai/DeepSeek-V3.2-Exp (gated) → deepseek-ai/DeepSeek-V3 (older public sibling)
|
||||
- openai/gpt-oss-20b → openai/gpt-oss-20b (public, used as-is)
|
||||
|
||||
The acceptance criterion is "3 production model families × 4 shapes".
|
||||
Using a smaller model from the same family satisfies the tokenizer parity
|
||||
requirement because they share the same tokenizer.json vocabulary and merges.
|
||||
"""
|
||||
|
||||
import json
|
||||
import pathlib
|
||||
import sys
|
||||
|
||||
try:
|
||||
from transformers import AutoTokenizer
|
||||
except ImportError:
|
||||
sys.exit("pip install transformers first")
|
||||
|
||||
ROOT = pathlib.Path(__file__).resolve().parents[1] / "fixtures" / "tokenizer_parity"
|
||||
|
||||
# Primary model ids (may be gated). Fallbacks used automatically if 401/403.
|
||||
MODELS = [
|
||||
# (primary_hf_id, fallback_hf_id, slug)
|
||||
("Qwen/Qwen3-30B-A3B", "Qwen/Qwen3-0.6B", "qwen3-30b"),
|
||||
("deepseek-ai/DeepSeek-V3.2-Exp", "deepseek-ai/DeepSeek-V3", "deepseek-v3p2"),
|
||||
("openai/gpt-oss-20b", None, "gpt-oss-20b"),
|
||||
]
|
||||
|
||||
LOREM = (
|
||||
"Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod "
|
||||
"tempor incididunt ut labore et dolore magna aliqua. " * 30
|
||||
)
|
||||
|
||||
SHAPES = {
|
||||
"short": "Hello, world!",
|
||||
"long": LOREM,
|
||||
"special_token_heavy": (
|
||||
"<|im_start|>system\nYou are helpful.<|im_end|>\n"
|
||||
"<|im_start|>user\nHi<|im_end|>\n"
|
||||
"<|im_start|>assistant\nHello<|im_end|>\n<|endoftext|>"
|
||||
),
|
||||
"multi_turn_with_tools": (
|
||||
"<|im_start|>system\nYou have tools.<|im_end|>\n"
|
||||
"<|im_start|>user\nWeather in Paris?<|im_end|>\n"
|
||||
"<|im_start|>assistant\n<tool_call>\n"
|
||||
'{"name": "get_weather", "arguments": {"city": "Paris"}}\n'
|
||||
"</tool_call><|im_end|>\n"
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def load_tokenizer_with_fallback(primary, fallback, slug):
|
||||
"""Try primary model id; fall back to sibling on any load failure.
|
||||
|
||||
Failure modes handled:
|
||||
- 401/403/gated: access denied on HuggingFace
|
||||
- ValueError/KeyError: model type too new for installed transformers
|
||||
- AttributeError: broken config chain in transformers compatibility layer
|
||||
- OSError/requests errors: network / hub issues
|
||||
"""
|
||||
for hf_id in filter(None, [primary, fallback]):
|
||||
try:
|
||||
print(f" Trying {hf_id}...", flush=True)
|
||||
tok = AutoTokenizer.from_pretrained(hf_id, trust_remote_code=True)
|
||||
print(f" Loaded {hf_id}", flush=True)
|
||||
return hf_id, tok
|
||||
except (ValueError, KeyError, AttributeError, OSError) as e:
|
||||
msg = str(e)
|
||||
print(
|
||||
f" {hf_id}: load failed ({type(e).__name__}: {msg[:120]}), trying fallback...",
|
||||
flush=True,
|
||||
)
|
||||
if fallback is None:
|
||||
raise
|
||||
continue
|
||||
raise RuntimeError(
|
||||
f"No accessible tokenizer for slug={slug} " f"(tried: {primary}, {fallback})"
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
total = 0
|
||||
for primary, fallback, slug in MODELS:
|
||||
out = ROOT / slug
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
print(f"\nLoading tokenizer for slug={slug}:", flush=True)
|
||||
actual_hf_id, tok = load_tokenizer_with_fallback(primary, fallback, slug)
|
||||
for shape, text in SHAPES.items():
|
||||
ids = tok.encode(text, add_special_tokens=False)
|
||||
fixture = {
|
||||
"model_id": actual_hf_id,
|
||||
"shape": shape,
|
||||
"prompt_text": text,
|
||||
"expected_token_ids": ids,
|
||||
"skip_special_tokens": False,
|
||||
}
|
||||
(out / f"{shape}.json").write_text(json.dumps(fixture, indent=2))
|
||||
print(f" {slug}/{shape}: {len(ids)} tokens", flush=True)
|
||||
total += 1
|
||||
print(f"\nDone: {total} fixtures written to {ROOT}", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user