[Feature] Add sampler custom logits processor

[hongpeng-guo]

Motivation

This PR tries to support custom logits processor registered by users, so users can easily implement their custom sampling methods without the need to change the sglang code.

Related Issue

#2291

Modifications

1. Introduce the common abstract interface CustomLogitProcessor as in sglang/srt/sampling/custom_logit_processor.py. This interface contains a (1) Callable function to process the logtis, and (2) from_str and to_str methods that any user defined subclass can be serialized and passed into the server using request.post(url, json=data) to the /generate endpoint.
2. Add custom_params as a field of SamplingParams
3. Add custom_logit_processor as a field to GenerateReqInput, TokenizedGenerateReqInput, Req, and SamplingBatchInfo

Examples

Frist define a dummy AddLogitProcessor

from sglang.srt.sampling.custom_logit_processor import CustomLogitProcessor

class AddLogitProcessor(CustomLogitProcessor):

    def __call__(self, logits, custom_param_list):
        import torch 
        assert logits.shape[0] == len(custom_param_list)
        key = 'arg1'

        merged_params = torch.tensor(
            [custom_param_list[i][key] for i in range(len(custom_param_list))], dtype=torch.float
        ).to(device=logits.device, non_blocking=True)
        return logits + merged_params

endpoint usage example:

import requests
    
url = "http://localhost:30000/generate"
data = {
    "text": "What is the capital of France?",
    "sampling_params": {        
        "custom_params": {
            "arg1": 5.0,
        },
    },
    "custom_logit_processor": AddLogitProcessor().to_str(),
}

response = requests.post(url, json=data)
print(response.json())

offline engine usage example

import sglang as sgl

llm = sgl.Engine(model_path="meta-llama/Meta-Llama-3.1-8B-Instruct")
prompt = "The president of the United States is"
sampling_params = {"temperature": 0.8, "top_p": 0.95, "custom_params": {"arg1": 5.0}}

output = llm.generate(
    prompt, 
    sampling_params,
    custom_logit_processor=AddLogitProcessor().to_str()
)
print(output['text'])

TODO

update the docs as suggested by @merrymercy

Checklist

• Format your code according to the Contributor Guide.
• Add unit tests as outlined in the Contributor Guide.
• Update documentation as needed, including docstrings or example tutorials.

[merrymercy]

Thanks for taking this. Can you add some end-to-end tests and examples?

[hongpeng-guo]

Thanks for taking this. Can you add some end-to-end tests and examples?

@merrymercy Thanks for taking a look on this. I am still trying to understand the appropriate layer for the user to register their customized_logit_processor. The goal is to enable customized_logit_processor functionality without requiring changes to the internal sglang codebase. To achieve this, the function registration should occur at the API layer. The customized_logit_processor_fn and custom_params will then be passed from the program driver to the internal modules, such as Sampler and SampleParams.

If the above seems correct, I will try to get it reviewable within this week.

[merrymercy]

Your proposal sounds good. We should be able to register this function with all interfaces.
Say, we should be able to support it with the native /generate api and the offline Engine API.

[hongpeng-guo]

There are a few unit test on performance and accuracy being a bit flaky. they seem not to be directly related to this PR, though.

[merrymercy]

We are almost there! Some final comments.

[hongpeng-guo]

We are almost there! Some final comments.

Thanks a lot for shepherding this! Just handled the comments. PTAL

[merrymercy]

we can even skip this for loop if custom logit process is not enabled by the server.

[merrymercy]

this seems not needed. If self.custom_logit_processor is None, why will it go to this function?

[hongpeng-guo]

Agreed, I can remove this or make it a simple assert.

[merrymercy]

@hongpeng-guo Thanks. It is merged.

Some follow-up items:

• Address the above two comments.
• Documentation. We can add some usage examples here https://github.com/sgl-project/sglang/blob/main/docs/references/sampling_params.md and it will be compiled as https://sgl-project.github.io/references/sampling_params.html