# /// 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()