sgl-project--sglang/python/sglang/test/kits/cache_hit_kit.py

import asyncio
import json
import time

import aiohttp
import requests

from sglang.benchmark.datasets.random import sample_random_requests
from sglang.benchmark.serving import RequestFuncOutput
from sglang.benchmark.utils import get_tokenizer, remove_prefix

AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60)


async def async_request_sglang_generate(
    payload,
    url,
    pbar=None,
):
    """Send a streaming request to the server and collect cache metrics.

    Returns a RequestFuncOutput with additional cached_tokens and output_ids attributes.
    """
    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
        headers = {}
        generated_text = ""
        all_output_ids = []
        ttft = 0.0
        st = time.perf_counter()
        most_recent_timestamp = st
        output = RequestFuncOutput()

        try:
            async with session.post(url=url, json=payload, headers=headers) as response:
                if response.status == 200:
                    prompt_tokens = 0
                    cached_tokens = 0

                    async for chunk_bytes in response.content:
                        chunk_bytes = chunk_bytes.strip()
                        if not chunk_bytes:
                            continue

                        chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ")
                        latency = time.perf_counter() - st

                        if chunk == "[DONE]":
                            pass
                        else:
                            data = json.loads(chunk)

                            # output_ids and text are always returned together
                            if data.get("output_ids"):
                                all_output_ids = data["output_ids"]
                                generated_text = data.get("text", "")
                                timestamp = time.perf_counter()

                                if ttft == 0.0:
                                    ttft = time.perf_counter() - st
                                    output.ttft = ttft
                                    prompt_tokens = (data.get("meta_info") or {}).get(
                                        "prompt_tokens", 0
                                    )
                                    cached_tokens = (data.get("meta_info") or {}).get(
                                        "cached_tokens", 0
                                    )
                                else:
                                    output.itl.append(timestamp - most_recent_timestamp)

                                most_recent_timestamp = timestamp

                    output.generated_text = generated_text
                    output.output_ids = all_output_ids
                    output.success = True
                    output.latency = latency
                    output.prompt_len = prompt_tokens
                    output.cached_tokens = cached_tokens
                    output.generated_len = len(output.itl) + 1
                else:
                    output.error = response.reason or ""
                    output.success = False
        except Exception as e:
            output.success = False
            output.error = str(e)
            print(f"Request failed: {e}")

    if pbar:
        pbar.update(1)
    return output


async def async_request_openai_chat_completions(
    payload,
    url,
    pbar=None,
):
    """Send a streaming request to an OpenAI-compatible /v1/chat/completions endpoint.

    Returns a RequestFuncOutput with the same dynamic attributes as
    async_request_sglang_generate (except output_ids, which is unavailable).
    """
    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
        generated_text = ""
        ttft = 0.0
        latency = 0.0
        st = time.perf_counter()
        most_recent_timestamp = st
        output = RequestFuncOutput()

        try:
            async with session.post(url=url, json=payload) as response:
                if response.status == 200:
                    prompt_tokens = 0
                    cached_tokens = 0
                    completion_tokens = 0

                    async for chunk_bytes in response.content:
                        chunk_bytes = chunk_bytes.strip()
                        if not chunk_bytes:
                            continue

                        chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ")
                        latency = time.perf_counter() - st

                        if chunk == "[DONE]":
                            pass
                        else:
                            data = json.loads(chunk)

                            # Streaming token chunks
                            if data.get("choices"):
                                raw_delta = data["choices"][0].get("delta")
                                text = raw_delta.get("content", "") if raw_delta else ""
                                if text:
                                    generated_text += text
                                    timestamp = time.perf_counter()

                                    if ttft == 0.0:
                                        ttft = time.perf_counter() - st
                                        output.ttft = ttft
                                    else:
                                        output.itl.append(
                                            timestamp - most_recent_timestamp
                                        )

                                    most_recent_timestamp = timestamp

                            # Final chunk with usage stats
                            usage = data.get("usage")
                            if usage:
                                prompt_tokens = usage.get("prompt_tokens", 0)
                                completion_tokens = usage.get("completion_tokens", 0)
                                details = usage.get("prompt_tokens_details", {}) or {}
                                cached_tokens = details.get("cached_tokens", 0)

                    output.generated_text = generated_text
                    output.output_ids = []  # Not available from OpenAI endpoint
                    output.success = True
                    output.latency = latency
                    output.prompt_len = prompt_tokens
                    output.cached_tokens = cached_tokens
                    output.generated_len = (
                        completion_tokens if completion_tokens else len(output.itl) + 1
                    )
                else:
                    output.error = response.reason or ""
                    output.success = False
        except Exception as e:
            output.success = False
            output.error = str(e)
            print(f"Request failed: {e}")

    if pbar:
        pbar.update(1)
    return output


def gen_payload_openai(messages, output_len, model):
    return {
        "model": model,
        "messages": messages,
        "max_tokens": output_len,
        "temperature": 0.0,
        "stream": True,
        "stream_options": {"include_usage": True},
    }


def gen_payload(input_ids, output_len, lora_path=""):
    return {
        "input_ids": input_ids,
        "sampling_params": {
            "temperature": 0.0,
            "max_new_tokens": output_len,
            "ignore_eos": True,
        },
        "stream": True,
        "stream_options": {"include_usage": True},
        "lora_path": lora_path,
        "return_logprob": False,
        "logprob_start_len": -1,
    }


async def _send_round(
    payloads,
    url,
    max_parallel,
):
    """Send a batch of payloads concurrently with concurrency limit."""
    semaphore = asyncio.Semaphore(max_parallel)

    async def _send_one(payload):
        async with semaphore:
            return await async_request_sglang_generate(payload, url)

    tasks = [asyncio.create_task(_send_one(p)) for p in payloads]
    return await asyncio.gather(*tasks)


def _get_page_size(base_url: str) -> int:
    """Query server for page_size used by radix cache."""
    try:
        resp = requests.get(f"{base_url}/server_info", timeout=10)
        resp.raise_for_status()
        info = resp.json()
        return info.get("page_size", 1)
    except Exception:
        return 1


def run_multiturn_cache_hit_test(
    base_url: str,
    model_path: str,
    num_clients: int = 8,
    num_rounds: int = 3,
    request_length: int = 256,
    output_length: int = 32,
    miss_tolerance: int = 1,
    sub_question_input_length: int = 0,
    lora_path: str = "",
    dataset_path: str = "",
    max_parallel: int = 64,
    seed: int = 1,
) -> dict:
    """Run a multi-turn workload and verify cache hit rate.

    Sends requests in round-barrier mode: all clients complete round i
    before round i+1 starts, ensuring deterministic cache state.

    The expected cache hit rate is self-computed from the workload structure:
    - Round 0: expected cached_tokens = 0 (cold start after flush)
    - Round r (r >= 1): each client's prefix from round r-1 should be cached,
      minus up to previous round's (prompt_len + decoding output - miss_tolerance) // page * page.

    Returns metrics dict with per-round and overall cache_hit_rate.
    """
    import random

    import numpy as np

    random.seed(seed)
    np.random.seed(seed)

    generate_url = f"{base_url}/generate"
    page_size = _get_page_size(base_url)

    # Flush cache for clean state
    requests.post(f"{base_url}/flush_cache")
    time.sleep(1)

    # Resolve sub-question length (0 means same as request_length)
    effective_sub_len = (
        sub_question_input_length if sub_question_input_length != 0 else request_length
    )

    # Sample initial prompts and sub-question prompts as token ids
    tokenizer = get_tokenizer(model_path)

    initial_inputs = sample_random_requests(
        input_len=request_length,
        output_len=output_length,
        num_prompts=num_clients,
        range_ratio=1.0,
        tokenizer=tokenizer,
        dataset_path=dataset_path,
        return_text=False,
    )
    # r.prompt is now List[int] when return_text=False
    initial_token_ids = [list(r.prompt) for r in initial_inputs]

    sub_question_inputs = sample_random_requests(
        input_len=effective_sub_len,
        output_len=output_length,
        num_prompts=num_clients * max(num_rounds - 1, 1),
        range_ratio=1.0,
        tokenizer=tokenizer,
        dataset_path=dataset_path,
        return_text=False,
    )
    sub_question_token_ids = [list(r.prompt) for r in sub_question_inputs]

    # Per-round metrics and per-client tracking for expected cache computation
    round_metrics = {
        i: {"prompt_len": [], "cached_tokens": [], "ttft": []}
        for i in range(num_rounds)
    }
    # Track the previous round's prompt_len per client to compute expected cache
    prev_prompt_lens = [0] * num_clients
    # histories now stores List[int] (token ids) for each client
    histories = [list(ids) for ids in initial_token_ids]
    sub_idx = 0

    for round_num in range(num_rounds):
        payloads = [gen_payload(h, output_length, lora_path) for h in histories]
        responses = asyncio.run(_send_round(payloads, generate_url, max_parallel))

        for i, resp in enumerate(responses):
            assert resp.success, f"Round {round_num}, client {i} failed: {resp.error}"

            round_metrics[round_num]["prompt_len"].append(resp.prompt_len)
            round_metrics[round_num]["cached_tokens"].append(resp.cached_tokens)
            round_metrics[round_num]["ttft"].append(resp.ttft)

            # Verify cache hit against expected value
            if round_num == 0:
                # Cold start: no cache expected
                expected_cached = 0
            else:
                # Previous round's prompt + output are in cache.
                # Radix cache aligns to page_size, so the last partial page
                # may not be cached.
                cacheable = prev_prompt_lens[i] + output_length - miss_tolerance
                expected_cached = (cacheable // page_size) * page_size

            msg = (
                f"Round {round_num}, client {i}: "
                f"cached_tokens={resp.cached_tokens}, "
                f"expected>={expected_cached} "
                f"(prev_prompt={prev_prompt_lens[i]}, "
                f"output={output_length}, page_size={page_size})"
            )

            print(msg)

            assert resp.cached_tokens >= expected_cached
            # Upper bound: cached tokens are a subset of the prompt, so they can
            # never exceed prompt_len. In PD disaggregation with decode radix
            # cache, the shared prefix was previously counted on both the prefill
            # and the decode node, making cached_tokens exceed prompt_len.
            assert resp.cached_tokens <= resp.prompt_len, (
                f"Round {round_num}, client {i}: cached_tokens="
                f"{resp.cached_tokens} exceeds prompt_len={resp.prompt_len} "
                f"(double-counted prefix across prefill/decode)"
            )

            # Record this round's prompt_len for next round's expected calc
            prev_prompt_lens[i] = resp.prompt_len

            # Accumulate history for next round using output_ids (token ids)
            histories[i].extend(resp.output_ids)
            if round_num < num_rounds - 1:
                histories[i].extend(sub_question_token_ids[sub_idx])
                sub_idx += 1

    # Compute per-round and overall cache hit rate
    total_prompt = 0
    total_cached = 0
    result = {"rounds": {}, "overall": {}}

    for r in range(num_rounds):
        rm = round_metrics[r]
        r_prompt = sum(rm["prompt_len"])
        r_cached = sum(rm["cached_tokens"])
        r_hit_rate = r_cached / r_prompt if r_prompt > 0 else 0.0
        r_avg_ttft = sum(rm["ttft"]) / len(rm["ttft"]) if rm["ttft"] else 0.0

        result["rounds"][f"round_{r}"] = {
            "cache_hit_rate": r_hit_rate,
            "average_ttft": r_avg_ttft,
            "total_prompt_tokens": r_prompt,
            "total_cached_tokens": r_cached,
            "request_count": len(rm["ttft"]),
        }

        total_prompt += r_prompt
        total_cached += r_cached

        print(
            f"  Round {r}: cache_hit_rate={r_hit_rate:.4f}, "
            f"avg_ttft={r_avg_ttft:.4f}s, "
            f"cached={r_cached}/{r_prompt} tokens"
        )

    overall_hit_rate = total_cached / total_prompt if total_prompt > 0 else 0.0
    result["overall"] = {
        "cache_hit_rate": overall_hit_rate,
        "total_prompt_tokens": total_prompt,
        "total_cached_tokens": total_cached,
    }
    print(f"  Overall cache_hit_rate={overall_hit_rate:.4f}")

    return result