How to reuse KV Cache in SGLang Frontend Language?

[mZhenz]

Hi, in SGLang Frontend Language, if I use s1 = s.fork(1)[0], can s1 reuse the kv cache of s?

I am using SGLang Frontend Language to implement blockwise parrallel decoding, and my implementation is as below.

However, I find that the runtime of draft and verify stage is same as prefill stage. It seems like it doesn't reuse the kv cache from prefill stage. How can I accelerate it? Is there any solution to reuse the kv cache?

Looking forward to your reply. Thanks a lot.

Environment

sglang=0.4.8

import time
from sglang import function, gen
from sglang import RuntimeEndpoint, set_default_backend

set_default_backend(RuntimeEndpoint(f"http://localhost:{port}"))
MASK_TOKEN = "<｜tool▁sep｜>"
MASK_TOKEN_NUM = 8
MAX_NEW_TOKENS = 64
STOP_TOKEN = ["\n", "<｜end▁of▁sentence｜>"]

@function
def blockwise_parallel_decoding(s, prompt):
    s += prompt
    start_time = time.time()
    s += gen(
        "prefill",
        temperature=0,
        max_tokens=0
    )
    meta_info = s.get_meta_info("prefill")
    end_time = time.time()
    print(f">>> prefill time: {end_time - start_time}")
    prompt_tokens = meta_info["prompt_tokens"]
    total_tokens = prompt_tokens
    
    for i in range(MAX_NEW_TOKENS):
        print("-"*32 + f"step{i}" + "-"*32)
        print(f"s.text(): {s.text()[len(prompt):]}")
        # draft
        draft_fork = s.fork(1)[0]
        draft_fork += MASK_TOKEN * MASK_TOKEN_NUM
        start_time = time.time()
        draft_fork += gen(
            "draft",
            temperature=0,
            max_tokens=0,
            return_logprob=True,
            top_logprobs_num=1,
            return_text_in_logprobs=True,
            logprob_start_len=total_tokens
        )
        meta_info = draft_fork.get_meta_info("draft")
        end_time = time.time()
        print(f">>> draft time: {end_time - start_time}")
        input_top_logprobs = meta_info["input_top_logprobs"]
        output_top_logprobs = meta_info["output_top_logprobs"]
        draft = [item[0][2] for item in input_top_logprobs[1:]] + [output_top_logprobs[0][0][2]]
        
        # verify
        verify_fork = s.fork(1)[0]
        verify_fork += "".join(draft)
        start_time = time.time()
        verify_fork += gen(
            "verify",
            temperature=0,
            max_tokens=0,
            return_logprob=True,
            top_logprobs_num=1,
            return_text_in_logprobs=True,
            logprob_start_len=total_tokens
        )
        meta_info = verify_fork.get_meta_info("verify")
        end_time = time.time()
        print(f">>> verify time: {end_time - start_time}")
        input_top_logprobs = meta_info["input_top_logprobs"]
        verify = [item[0][2] for item in input_top_logprobs[1:]]
        
        # determine how many tokens to accept
        start_time = time.time()
        acceptance = ""
        for d, v in zip(draft, verify):
            acceptance += v
            total_tokens += 1
            if v != d: break
        s += acceptance
        
        if any(stop in acceptance for stop in STOP_TOKEN):
            break
        if (total_tokens - prompt_tokens) >= MAX_NEW_TOKENS:
            break
        end_time = time.time()
        print(f">>> acception: {end_time - start_time}")