动手 4：写自定义 DSL 程序——检索 + 多分支推理 + 投票

用户问题
  │
  ▼
1. 检索：从向量库拿 3 篇相关文档
  │
  ▼
2. 多分支推理：3 个并行分支，各自基于 3 篇文档生成答案
  │
  ▼
3. 投票：让 LLM 在 3 个候选答案里选最准确的
  │
  ▼
4. 结构化输出：{"answer": "...", "sources": ["..."], "confidence": 0.x}

import sglang as sgl
import json

sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000"))

# 假设这是检索结果（实际从向量库来）
def retrieve(question: str) -> list[str]:
    return [
        f"Document A about {question}: ...",
        f"Document B about {question}: ...",
        f"Document C about {question}: ...",
    ]


SYSTEM = (
    "You are a knowledge-base QA assistant. "
    "Answer based on the provided documents. "
    "If documents don't contain the answer, say 'I don't know'."
)


OUTPUT_SCHEMA = {
    "type": "object",
    "properties": {
        "answer": {"type": "string"},
        "sources": {"type": "array", "items": {"type": "string"}},
        "confidence": {"type": "number", "minimum": 0, "maximum": 1}
    },
    "required": ["answer", "sources", "confidence"]
}


@sgl.function
def kb_qa(s, question, docs):
    # 公共前缀（所有分支共享 → cache 命中）
    s += sgl.system(SYSTEM)
    s += sgl.user(f"Question: {question}")

    # 三分支推理
    forks = s.fork(3)
    for i, f in enumerate(forks):
        f += "\nDocument:\n" + docs[i]
        f += "\nAnswer based on this document only:\n"
        f += sgl.assistant_begin()
        f += sgl.gen(f"draft_{i}", max_tokens=300, temperature=0.3)
        f += sgl.assistant_end()

    drafts = [forks[i][f"draft_{i}"] for i in range(3)]

    # 投票：让 LLM 选最准确的
    s += "\n\nThree candidate answers:\n"
    for i, d in enumerate(drafts):
        s += f"\n[Candidate {i}]\n{d}\n"
    s += "\nPick the most accurate candidate and output structured JSON.\n"
    s += sgl.assistant_begin()
    s += sgl.gen(
        "final",
        max_tokens=500,
        temperature=0.0,
        json_schema=json.dumps(OUTPUT_SCHEMA),
    )
    s += sgl.assistant_end()


def serve(question: str):
    docs = retrieve(question)
    state = kb_qa.run(question=question, docs=docs)
    return json.loads(state["final"])


# 跑一次
result = serve("RadixAttention 解决了什么问题？")
print(json.dumps(result, indent=2, ensure_ascii=False))

请求生命周期（一次 serve）:

1) Scheduler 收到一个根请求 (system + user + question)
   → match_prefix 命中 SYSTEM（如果之前发过）
   → prefill 末尾差异
   → 出第一个 token，存到 root 子节点

2) Fork 3 个子分支:
   分支 0: prefix = root → +"\nDocument: A\n..."
   分支 1: prefix = root → +"\nDocument: B\n..."
   分支 2: prefix = root → +"\nDocument: C\n..."
   →三个分支共享 root 的 KV（不复制）

3) 三个分支并行 prefill + decode（cascade attention 高效）
   → 各自产出 draft_i

4) 收集 drafts 拼回主 state
   → 新的 prefix = root + drafts → match_prefix 命中 root
   → 新的 prompt 是 "三个候选 + 选最优" → prefill 末尾
   → gen("final", json_schema=...) 触发受约束输出
   → 出 JSON 答案

state = kb_qa.run(question=question, docs=docs, stream=True)
for chunk in state.text_iter("final"):
    print(chunk, end="", flush=True)

n_branches = 3 if is_simple(question) else 5

s += sgl.select("best_idx", choices=["0", "1", "2"])
# 然后取 drafts[int(state["best_idx"])]

import time

# Cold start
t0 = time.time()
serve("第一个问题")
cold = time.time() - t0

# Warm
t0 = time.time()
serve("第一个问题")           # 100% cache
warm_same = time.time() - t0

t0 = time.time()
serve("不同的第二个问题")
warm_diff = time.time() - t0

print(f"Cold: {cold:.2f}s, Warm same: {warm_same:.2f}s, Warm diff: {warm_diff:.2f}s")

curl http://localhost:30000/metrics | grep cache_hit_rate