SGLang 结构化生成引擎

import sglang as sgl
print(sgl.__version__)

@sgl.function
def simple_qa(s, question):
    s += question
    s += sgl.gen("answer", max_tokens=256, temperature=0.7)

@sgl.function
def sentiment_analysis(s, text):
    s += "分析以下文本的情感倾向：\n"
    s += text + "\n"
    s += "情感倾向："
    s += sgl.select("sentiment", choices=["正面", "负面", "中性"])

@sgl.function
def image_qa(s, image_path, question):
    s += sgl.image(image_path)
    s += question
    s += sgl.gen("answer", max_tokens=256)

import sglang as sgl

@sgl.function
def complex_workflow(s, topic):
    # 步骤1：生成主题大纲
    s += "请为以下主题生成大纲：" + topic + "\n"
    s += "大纲："
    s += sgl.gen("outline", max_tokens=256, temperature=0.7)

    # 步骤2：判断大纲质量
    s += "\n这份大纲是否足够详细？"
    s += sgl.select("quality", choices=["详细", "简略", "需要改进"])

    # 步骤3：根据判断结果生成详细内容
    s += "\n请根据大纲展开详细内容：\n"
    s += sgl.gen("detail", max_tokens=512, temperature=0.8)

from sglang import Runtime

# 基本创建
runtime = Runtime(
    model_path="meta-llama/Llama-2-7b-hf",  # 模型路径
)

# 完整参数创建
runtime = Runtime(
    model_path="meta-llama/Llama-2-7b-hf",
    tp_size=1,                    # 张量并行大小
    mem_fraction_static=0.88,     # 静态分配的显存比例
    base_gpu_id=0,                # 起始GPU ID
    enable_radix_cache=True,      # 启用RadixAttention（默认True）
    enable_torch_compile=False,   # 启用TorchCompile优化
)

import sglang as sgl

# 初始化运行时
runtime = sgl.Runtime(model_path="meta-llama/Llama-2-7b-hf")

# 定义生成函数
@sgl.function
def qa(s, question):
    s += "问题：" + question + "\n答案："
    s += sgl.gen("answer", max_tokens=128, temperature=0.7)

# 批量运行
questions = ["什么是AI？", "什么是深度学习？", "什么是NLP？"]
states = qa.run_batch(
    [{"question": q} for q in questions],
    runtime=runtime,
    num_threads="auto",  # 自动选择并发线程数
)

# 获取结果
for state in states:
    print(state["answer"])

# 清理缓存
runtime.flush_cache()

runtime.flush_cache()  # 清空RadixAttention缓存的所有KV cache

@sgl.function
def extract_info(s, text):
    s += "从以下文本中提取日期：\n"
    s += text + "\n"
    s += "日期："
    # 约束输出为日期格式
    s += sgl.gen("date", regex=r"\d{4}-\d{2}-\d{2}")

@sgl.function
def extract_number(s, text):
    s += "提取价格：\n"
    s += text + "\n价格："
    # 约束输出为数字格式
    s += sgl.gen("price", regex=r"\d+\.?\d*")

import json

@sgl.function
def json_output(s, query):
    s += "请以JSON格式回答：\n"
    s += query + "\n"
    s += "```json\n"
    # 使用regex约束JSON格式
    s += sgl.gen(
        "json_result",
        max_tokens=512,
        stop="```",
        regex=r'\{[^}]*\}',  # 约束为JSON对象
    )

# 使用示例
runtime = sgl.Runtime(model_path="meta-llama/Llama-2-7b-hf")
state = json_output.run(
    query="列出三种编程语言及其特点",
    runtime=runtime,
)
result = json.loads(state["json_result"])

@sgl.function
def classify(s, text):
    s += "分类以下文本：\n" + text + "\n类别："
    s += sgl.select("category", choices=["科技", "体育", "娱乐", "政治"])

@sgl.function
def yes_no_qa(s, question):
    s += question + "\n答案："
    s += sgl.select("answer", choices=["是", "否"])

import openai

client = openai.OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="not-needed",
)

# Chat Completions
response = client.chat.completions.create(
    model="meta-llama/Llama-2-7b-hf",
    messages=[
        {"role": "system", "content": "你是一个有帮助的AI助手。"},
        {"role": "user", "content": "解释什么是深度学习"},
    ],
    max_tokens=256,
    temperature=0.7,
)
print(response.choices[0].message.content)

# Completions
response = client.completions.create(
    model="meta-llama/Llama-2-7b-hf",
    prompt="深度学习是",
    max_tokens=128,
)
print(response.choices[0].text)

import sglang as sgl

@sgl.function
def multi_turn_chat(s, messages):
    """多轮对话，自动复用历史KV cache"""
    for msg in messages:
        if msg["role"] == "user":
            s += "用户：" + msg["content"] + "\n"
        elif msg["role"] == "assistant":
            s += "助手：" + msg["content"] + "\n"
    s += "助手："
    s += sgl.gen("response", max_tokens=256, temperature=0.7)

runtime = sgl.Runtime(model_path="meta-llama/Llama-2-7b-hf")

# 第一轮对话
messages_1 = [{"role": "user", "content": "什么是机器学习？"}]
state_1 = multi_turn_chat.run(messages=messages_1, runtime=runtime)
print(state_1["response"])

# 第二轮对话（复用第一轮的KV cache）
messages_2 = messages_1 + [
    {"role": "assistant", "content": state_1["response"]},
    {"role": "user", "content": "它和深度学习有什么区别？"},
]
state_2 = multi_turn_chat.run(messages=messages_2, runtime=runtime)
print(state_2["response"])
# RadixAttention自动识别前两轮对话的公共前缀并复用KV cache

import sglang as sgl
import json

@sgl.function
def agent_workflow(s, query):
    """ReAct风格的Agent工作流"""
    # 第1步：思考
    s += "查询：" + query + "\n"
    s += "思考：我需要先分析这个问题。\n"
    s += "思考："
    s += sgl.gen("thought_1", max_tokens=128, temperature=0.7)

    # 第2步：决定行动
    s += "\n行动类型："
    s += sgl.select("action_type", choices=["搜索", "计算", "回答"])

    # 第3步：根据行动类型分支
    if s["action_type"] == "回答":
        s += "\n最终答案："
        s += sgl.gen("answer", max_tokens=256, temperature=0.5)
    else:
        s += "\n行动内容："
        s += sgl.gen("action_content", max_tokens=128)
        s += "\n观察结果：[模拟观察]\n最终答案："
        s += sgl.gen("answer", max_tokens=256, temperature=0.5)

runtime = sgl.Runtime(model_path="meta-llama/Llama-2-7b-hf")

# 运行Agent（多个Agent调用共享相同的前缀结构）
queries = [
    "法国的首都人口是多少？",
    "Python的GIL是什么？",
    "量子纠缠的原理是什么？",
]
states = agent_workflow.run_batch(
    [{"query": q} for q in queries],
    runtime=runtime,
)
for state in states:
    print(f"答案: {state['answer']}")

import sglang as sgl
import json

@sgl.function
def extract_structured_data(s, text):
    """从非结构化文本中提取结构化数据"""
    s += "从以下文本中提取信息：\n"
    s += text + "\n\n"

    # 使用约束解码确保输出格式
    s += "姓名："
    s += sgl.gen("name", max_tokens=20, stop="\n")
    s += "\n年龄："
    s += sgl.gen("age", regex=r"\d{1,3}", max_tokens=5)
    s += "\n职业："
    s += sgl.gen("occupation", max_tokens=30, stop="\n")
    s += "\n城市："
    s += sgl.gen("city", max_tokens=20, stop="\n")
    s += "\n情感："
    s += sgl.select("sentiment", choices=["积极", "消极", "中性"])

runtime = sgl.Runtime(model_path="meta-llama/Llama-2-7b-hf")

texts = [
    "张三今年28岁，是一名软件工程师，居住在北京，对AI技术充满热情。",
    "李四今年35岁，是一名医生，居住在上海，最近工作压力很大。",
]

states = extract_structured_data.run_batch(
    [{"text": t} for t in texts],
    runtime=runtime,
)

for state in states:
    print(json.dumps({
        "name": state["name"],
        "age": state["age"],
        "occupation": state["occupation"],
        "city": state["city"],
        "sentiment": state["sentiment"],
    }, ensure_ascii=False, indent=2))

import sglang as sgl
from collections import Counter

@sgl.function
def math_reasoning(s, problem):
    """数学推理：生成多条推理路径"""
    s += "问题：" + problem + "\n"
    s += "请逐步推理：\n"
    s += sgl.gen("reasoning", max_tokens=512, temperature=0.7)
    s += "\n因此，答案是："
    s += sgl.gen("answer", max_tokens=20, stop="\n")

runtime = sgl.Runtime(model_path="meta-llama/Llama-2-7b-hf")

problem = "一个农场有鸡和兔子共35个头，94只脚，问鸡和兔子各多少只？"

# 生成多条推理路径
states = math_reasoning.run_batch(
    [{"problem": problem} for _ in range(8)],  # 8条推理路径
    runtime=runtime,
    num_threads=8,
)

# 投票选择最常见答案
answers = [s["answer"].strip() for s in states]
most_common = Counter(answers).most_common(1)[0]
print(f"多数投票答案: {most_common[0]} (出现{most_common[1]}次)")

import sglang as sgl

@sgl.function
def llm_judge(s, question, answer_a, answer_b):
    """使用LLM作为裁判评估两个回答"""
    s += "请评估以下两个回答的质量：\n\n"
    s += "问题：" + question + "\n\n"
    s += "回答A：" + answer_a + "\n\n"
    s += "回答B：" + answer_b + "\n\n"
    s += "评估标准：准确性、完整性、清晰度\n\n"
    s += "更好的回答是："
    s += sgl.select("winner", choices=["A", "B", "两者相当"])

runtime = sgl.Runtime(model_path="meta-llama/Llama-2-7b-hf")

question = "什么是量子纠缠？"
answer_a = "量子纠缠是两个粒子之间的关联..."
answer_b = "量子纠缠指两个量子系统之间存在非经典的相关性..."

state = llm_judge.run(
    question=question,
    answer_a=answer_a,
    answer_b=answer_b,
    runtime=runtime,
)
print(f"裁判结果: {state['winner']}")

# 构建阶段：创建SGLang程序
@sgl.function
def my_program(s, x):
    s += "问题：" + x + "\n"
    s += "思考："
    s += sgl.gen("thought", max_tokens=128)
    s += "\n答案："
    s += sgl.gen("answer", max_tokens=64)

# 执行阶段：运行程序
state = my_program.run(x="什么是AI？", runtime=runtime)

# RadixAttention核心逻辑（简化版）
class RadixAttention:
    def __init__(self):
        self.tree = RadixTree()       # 基数树
        self.cache = {}               # token_hash -> KV cache blocks
        self.lru_queue = []           # LRU淘汰队列

    def match_prefix(self, token_ids):
        """在基数树中查找最长公共前缀"""
        node = self.tree.root
        matched_length = 0
        for i, token in enumerate(token_ids):
            if token in node.children:
                node = node.children[token]
                matched_length = i + 1
            else:
                break
        return matched_length, node

    def insert(self, token_ids, kv_cache_blocks):
        """将新序列的KV cache插入基数树"""
        matched_len, matched_node = self.match_prefix(token_ids)

        # 复用已匹配部分的KV cache
        reused_blocks = self._get_blocks(matched_node)

        # 只为新增部分分配和计算KV cache
        new_tokens = token_ids[matched_len:]
        new_blocks = self._allocate_and_compute(new_tokens)

        # 插入基数树
        self.tree.insert(matched_node, new_tokens, new_blocks)

        return reused_blocks, new_blocks

    def evict(self, num_blocks_needed):
        """LRU策略淘汰KV cache"""
        while num_blocks_needed > 0 and self.lru_queue:
            victim = self.lru_queue.pop(0)  # 淘汰最久未使用
            if victim.is_leaf():
                self._free_blocks(victim.blocks)
                self.tree.remove(victim)
                num_blocks_needed -= len(victim.blocks)

# SGLang调度器核心逻辑（简化版）
class SGLangScheduler:
    def schedule(self, waiting_queue, running_queue):
        """每次迭代的调度"""
        # 1. 处理running队列（继续解码）
        for req in running_queue:
            if req.is_finished():
                running_queue.remove(req)
                # 注意：完成的请求的KV cache保留在基数树中
                # 以便后续请求复用
            else:
                req.decode_step()

        # 2. 从waiting队列调度新请求
        for req in waiting_queue:
            # 检查RadixAttention中的前缀匹配
            matched_len = self.radix_attn.match_prefix(req.token_ids)

            # 只需为新增部分分配显存
            new_tokens = len(req.token_ids) - matched_len
            if self.has_memory_for(new_tokens):
                req.set_prefix_match(matched_len)
                running_queue.append(req)
                waiting_queue.remove(req)

# 推荐场景：多个请求共享前缀
# - 多轮对话：历史消息作为共享前缀
# - RAG：系统提示和检索上下文作为共享前缀
# - Agent工作流：工具描述和指令作为共享前缀
runtime = sgl.Runtime(
    model_path="model-name",
    enable_radix_cache=True,  # 默认True，建议保持
)

# 不推荐场景：每个请求完全独立，无共享前缀
# 此时RadixAttention会增加少量查找开销
# 可以禁用以获得更好的单请求性能
runtime = sgl.Runtime(
    model_path="model-name",
    enable_radix_cache=False,  # 完全独立请求时
)

# 使用run_batch进行批量推理，而不是循环调用run
questions = ["问题1", "问题2", "问题3", ...]

# 好的做法：批量推理
states = qa.run_batch(
    [{"question": q} for q in questions],
    runtime=runtime,
    num_threads="auto",  # 自动选择最优并发数
)

# 不好的做法：逐个推理
states = []
for q in questions:
    state = qa.run(question=q, runtime=runtime)
    states.append(state)

@sgl.function
def constrained_output(s, query):
    # 1. 使用select代替gen+stop进行分类
    s += "情感："
    s += sgl.select("sentiment", choices=["正面", "负面", "中性"])
    # 比以下方式更高效、更准确：
    # s += sgl.gen("sentiment", max_tokens=5, stop=["正面", "负面", "中性"])

    # 2. 使用regex确保格式正确
    s += "\n日期："
    s += sgl.gen("date", regex=r"\d{4}-\d{2}-\d{2}")
    # 比依赖模型自行格式化更可靠

    # 3. 约束范围不要过于严格
    # 好的做法：合理的正则范围
    s += "\n金额："
    s += sgl.gen("amount", regex=r"\d+\.?\d{0,2}")

    # 不好的做法：过于严格的约束可能导致模型困惑
    # s += sgl.gen("amount", regex=r"12345\.67")  # 太严格

# 调整静态显存比例
runtime = sgl.Runtime(
    model_path="model-name",
    mem_fraction_static=0.88,  # 默认0.88
    # 如果OOM，可以降低此值（但会降低KV cache容量）
    # 如果显存充足，可以适当提高
)

# 定期清理缓存（通常不需要）
# RadixAttention会自动管理缓存淘汰
runtime.flush_cache()  # 仅在切换完全不同的任务时使用