File size: 2,288 Bytes
1470c84 e27e171 1470c84 e27e171 1470c84 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | # /// script
# requires-python = ">=3.12,<3.13"
# dependencies = [
# "sglang==0.5.19",
# "torch==2.13.0+cu130",
# ]
# [tool.uv]
# index-strategy = "unsafe-best-match"
# prerelease = "allow"
# [[tool.uv.index]]
# url = "https://download.pytorch.org/whl/cu130"
# ///
"""Serve Ripple + DFlash2 on two B200s. Accepts standard SGLang CLI flags."""
import os
import shlex
import sys
from pathlib import Path
sys.dont_write_bytecode = True
# SGLang sizes the draft's positional cache for the target context, even with
# windowed draft attention. Allow its 128Ki position limit to grow to 1Mi;
# --speculative-draft-window-size still limits draft attention to 16Ki tokens.
os.environ.setdefault("SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN", "1")
# Top-level registration also runs in SGLang's spawned worker processes.
from sglang.srt.models.registry import ModelRegistry
from model import RicursiveForCausalLM
import ripple_format
ModelRegistry.models["RicursiveForCausalLM"] = RicursiveForCausalLM
ripple_format.register()
DEFAULT_ARGS = shlex.split("""
--host 127.0.0.1 --served-model-name ripple --trust-remote-code
--dtype bfloat16 --kv-cache-dtype bfloat16 --tp-size 2
--attention-backend trtllm_mha --page-size 64
--speculative-algorithm DFLASH --speculative-draft-model-quantization unquant
--speculative-num-draft-tokens 16 --speculative-draft-window-size 16384
--speculative-draft-attention-backend triton
--context-length 1048576 --mem-fraction-static 0.80 --max-running-requests 16
--cuda-graph-backend-prefill disabled
--tool-call-parser ripple --reasoning-parser ripple
""")
DEFAULT_ARGS += [
"--chat-template",
str(Path(__file__).resolve().parents[2] / "chat_template.jinja"),
]
def main():
from sglang.srt.entrypoints.http_server import launch_server
from sglang.srt.server_args import prepare_server_args
from sglang.srt.utils import kill_process_tree
args = prepare_server_args(DEFAULT_ARGS + sys.argv[1:])
if args.dtype != "bfloat16" or args.kv_cache_dtype != "bfloat16":
raise ValueError("Only BF16 weights and KV cache are supported")
try:
launch_server(args)
finally:
kill_process_tree(os.getpid(), include_parent=False)
if __name__ == "__main__":
main()
|