Set server context to 1Mi with a 16Ki draft window

e27e171 verified about 12 hours ago

2.29 kB

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

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