Summary of Changes

Hello @Ubospica, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the SGLang system by integrating support for constrained decoding with Speculative Decoding V2 (EAGLE). The core purpose is to enable the simultaneous use of both advanced features, allowing for more controlled and structured text generation while retaining the performance benefits of speculative decoding. The changes involve modifying the token sampling process to incorporate grammar-based vocabulary masks and adding comprehensive tests to ensure the correctness and compatibility of these combined functionalities across different speculative decoding versions and constraint types.

Highlights

• Constrained Decoding Integration: Implemented the ability to apply vocabulary masks during speculative decoding (Spec V2) to enforce grammar constraints, ensuring generated tokens adhere to specified rules.
• Grammar Data Preparation: Added logic within the verification process to prepare grammar-related data on the CPU and dynamically generate a vocab_mask based on active grammar rules for constrained decoding.
• Comprehensive Test Coverage: Introduced new unit tests to validate the combined functionality of speculative decoding (both EAGLE v1 and v2) with JSON schema and regex-based constrained decoding, ensuring robustness.

Thanks, I test your impl use Llama3.1-8b-Instruct and eagle model. When set export SGLANG_ENABLE_SPEC_V2=0, the response satisfies r"^user@example\.com$" . When set export SGLANG_ENABLE_SPEC_V2=1, the response is use the following information to generate an email address. I test use the following prompt and sampling parameter.

    payload = {
        "text": prompt,
        "sampling_params": {
            "temperature": 0.5,
            "max_new_tokens": 128,
            "n": 3,
            "regex": regex,
        },
        "stream": False,
        "return_logprob": False,
        "top_logprobs_num": 0,
        "logprob_start_len": 0,
    }

/tag-run-ci-label try again

Thanks, I test your impl use Llama3.1-8b-Instruct and eagle model. When set export SGLANG_ENABLE_SPEC_V2=0, the response satisfies r"^user@example\.com$" . When set export SGLANG_ENABLE_SPEC_V2=1, the response is use the following information to generate an email address. I test use the following prompt and sampling parameter.

    payload = {
        "text": prompt,
        "sampling_params": {
            "temperature": 0.5,
            "max_new_tokens": 128,
            "n": 3,
            "regex": regex,
        },
        "stream": False,
        "return_logprob": False,
        "top_logprobs_num": 0,
        "logprob_start_len": 0,
    }

@jiapingW Thanks for testing! I will further fix this.

Thanks, I test your impl use Llama3.1-8b-Instruct and eagle model. When set export SGLANG_ENABLE_SPEC_V2=0, the response satisfies r"^user@example\.com$" . When set export SGLANG_ENABLE_SPEC_V2=1, the response is use the following information to generate an email address. I test use the following prompt and sampling parameter.

    payload = {
        "text": prompt,
        "sampling_params": {
            "temperature": 0.5,
            "max_new_tokens": 128,
            "n": 3,
            "regex": regex,
        },
        "stream": False,
        "return_logprob": False,
        "top_logprobs_num": 0,
        "logprob_start_len": 0,
    }

@jiapingW Thanks for testing! I will further fix this.
I'm also fixing it. You can check here first: the EAGLEWorkerV2 verifyfunction's input batch is a ModelWorkerBatchclass. It doesn't have has_grammar attribute. If you have fixed it, I'll help to test it.

My main finding is that because of the spec v2 overlap, the grammar is not updated immediately after prefilling, but only after the first decode. This results in the grammar being None during the first decode, leading to duplicate data generation, which also affects subsequent results.

I impl a runable version with no polish in https://github.com/sgl-project/sglang/pull/13441/files. However, I haven't conducted any testing or performance analysis yet. I test use the following code. Its result is OK.

import json
import re
import sys
import requests
import argparse

def regex_match(text, pattern):
    return re.fullmatch(pattern, text) is not None

def send_sglang_request(
    base_url,
    prompt,
    regex,
    n=1,
    temperature=0,
    max_new_tokens=128,
):
    # if n > 1 and temperature == 0:
    #     temperature = 0.5
    #     print(f"Warning: n > 1 but temperature is 0. Automatically setting temperature to {temperature} for diverse outputs.")
    temperature = 0.5
    payload = {
        "text": prompt,
        "sampling_params": {
            "temperature": temperature,
            "max_new_tokens": max_new_tokens,
            "n": n,
            "regex": regex,
        },
        "stream": False,
        "return_logprob": False,
        "top_logprobs_num": 0,
        "logprob_start_len": 0,
    }

    print("--- Sending Request ---")
    print("URL:", base_url + "/generate")
    print("Payload:")
    print(json.dumps(payload, indent=2))
    print("-" * 25)

    try:
        response = requests.post(base_url + "/generate", json=payload)
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        print(f"Error: Request failed: {e}")
        sys.exit(1)

    ret = response.json()

    print("\n--- Received Response ---")
    print("Raw JSON Response:")
    # print(ret)
    # for i,item in enumerate(ret):
    #     print(f"response_{i}:",item["text"])
    print("-" * 25)

    if isinstance(ret, dict):
        results_list = [ret]
    elif isinstance(ret, list):
        results_list = ret
    else:
        print(f"Error: Expected response to be a list or dict, but got {type(ret)}")
        sys.exit(1)

    print("\n--- Validating Results ---")
    all_passed = True
    if not results_list:
        print("Error: Response list is empty.")
        all_passed = False

    for i, item in enumerate(results_list):
        text = item.get("text", "").strip()
        print(f"Result {i+1}/{len(results_list)}: '{text}'")

        if not text:
            print("  -> FAIL: Generated text is empty.")
            all_passed = False
            continue

        if not regex_match(text, regex):
            print(f"  -> FAIL: Text does not match regex pattern '{regex}'.")
            all_passed = False
        else:
            print("  -> PASS: Text matches regex.")
    
    print("-" * 25)
    return all_passed


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test SGLang server with regex-constrained generation.")
    parser.add_argument("--port", type=int, default=21000, help="Port of the SGLang server (default: 30000)")
    args = parser.parse_args()
    test_prompt = "Generate an email address:"
    test_regex = r"^user@example\.com$"
    num_generations = 3

    print(f"Running test: Generate {num_generations} email(s) with regex '{test_regex}'")
    print("=" * 100)

    SGLANG_BASE_URL = f"http://127.0.0.1:{args.port}"

    success = send_sglang_request(
        base_url=SGLANG_BASE_URL,
        prompt=test_prompt,
        regex=test_regex,
        n=num_generations,
    )

    print("\n--- Final Result ---")
    if success:
        print("✅ All generated texts passed validation.")
        sys.exit(0)
    else:
        print("❌ One or more generated texts failed validation.")
        sys.exit(1)

My design is as follows:

Take question: "Generate an email address:" and grammar: "^user@example\.com$" as an example.

The original Spec V2's overlap design handles the process as follows:

1. After the prefill run_batch finishes, it produces the "use" token. At this point, the grammar is not processed because last_batch is None.
2. After the first decode-phase run_batch finishes, it produces the user token. Since the grammar has not been updated, it still constructs the vocab mask based on the initial user. This results in repeated generation of matching tokens, producing outputs similar to useuseruser@eexamplele.com.

Therefore, my design idea is: place the grammar update operation after run_batch, and separate process_batch_result and grammar processing as much as possible. This way, process_batch_result can handle the output normally, while the grammar value is always up-to-date and won't affect subsequent token generation.
The revised processing steps are:

• 1. After the prefill run_batch finishes, it produces the "use" token. At this point, update the grammar value to use.
• 2. After entering the first decode phase run_batch finishes, it produces the r token. After completion, update the grammar value to user, and simultaneously execute process_batch_result to handle the batch result and output the result token use.
• .......
• n. After entering the n-th decode phase run_batch finishes, it produces the .com<|im_end|> token. After completion, update the grammar value to user@example.com<im_end>, and execute process_batch_result to process the batch result and output the result token ple.
• n+1. After entering the (n+1)-th round of the decode phase, run_batch. The input is <im_end>. After completion, token number 151935 is generated. The grammar is not updated. Simultaneously, process_batch_result is executed to process the batch result and output the result to .com.

Need to fix

The final round should not need to execute this special run_batch; it can be skipped directly by checking whether the grammar has terminated.

@Ubospica Can help review the impl?

@Ubospica Can help review the impl?

Certainly. I will check that out. Thanks for the reference!

/tag-and-rerun-ci

Now all the tests have passed. cc @hnyls2002

/tag-and-rerun-ci