Perform Batch Tokenization.

[sundar24295s]

Motivation

• The current tokenizer_manager processes batch workloads by sequentially tokenizing each prompt before sending it to the DP. This creates overhead in the tokenization phase.
• We can improve tokenization performance by leveraging the batch tokenization capabilities of fast tokenizers to process multiple prompts simultaneously before distributing them across DPs.

Modifications

• This PR introduces a batch tokenization option in TokenizerManager, controlled by server_args.enable_tokenizer_batch_encode.
  • When enabled, it tokenizes all text inputs in a single pass for generation requests, while other input types—such as input_ids, input_embeds, or image data—continue to follow the existing sequential process.
  • This optimization is particularly beneficial for prefill-heavy workloads with smaller batch sizes
• Also, added couple of benchmarking scripts:
  • To benchmark the tokenizer with batch prompts.
  • To benchmark sending batch requests to the /generate endpoint.

Future Work

• The current implementation tokenizes the entire batch at once which is suitable for usecases sending prompts in smaller batch sizes. But, which may not scale well for very large batch sizes (e.g., 1000).
• Future improvements could include:
  • Splitting large batches into manageable chunks, tokenize the chunk and send them to the DPs
• These enhancements will be explored in future iterations to support a wider range of use cases

Benchmarks

• All benchmarks are performed on H100s.

Batch Tokenization

(venv) jobuser [ ~/sglang ]$ python3.10 benchmark/benchmark_batch/benchmark_tokenizer.py
Tokenizer Benchmark: Sequential vs Batch Processing
------------------------------------------------------------
Tokenizer: /shared/public/sharing/fait360brew/training/models/meta-llama/Llama-3.2-3B
Tokens per prompt: 20000
Number of runs per batch size: 5
------------------------------------------------------------
Generating 8 random prompts with 20000 tokens each...
  Prompt 0: 20905 tokens
  Prompt 1: 20867 tokens
  Prompt 2: 20889 tokens
  Prompt 3: 20882 tokens
  Prompt 4: 20786 tokens
  Prompt 5: 20891 tokens
  Prompt 6: 20876 tokens
  Prompt 7: 20835 tokens

Running benchmark...
.
.
.
============================================================
SUMMARY OF RESULTS
============================================================
Batch Size Sequential (ms)    Batch (ms)         Speedup   
------------------------------------------------------------
1          33.23 ms         33.15 ms         1.00x
2          67.28 ms         39.67 ms         1.70x
4          159.67 ms         57.98 ms         2.75x
8          351.50 ms         67.81 ms         5.18x

Bechmark Batch prefill

• The following benchmark is to paint a picture on how much overhead we can save for a batched request if we perform batch tokenization.
• Launch Server Command

  (venv) jobuser [ ~/sglang ]$ python -m sglang.launch_server --model-path /models/meta-llama/Llama-3.2-3B --port 30000 --host 0.0.0.0 --disable-radix-cache --disable-cuda-graph --max-prefill-tokens 131072 --chunked-prefill-size 131072 --tp 1 --dp 8

Results using exisitng Sequential Tokenization

(venv) jobuser [ ~/sglang ]$ python3.10 benchmark/benchmark_batch/benchmark_batch.py 
.
.
.
Generated 480 prompts with 32000 tokens each, grouped into 60 requests of 8 prompts.

Starting benchmark: NUM_TOKENS=32000, BATCH_SIZE=8, NUM_REQUESTS=60

[Request] Sending request 1/10 with 8 prompts at 1744061155472
.
.
.

Benchmark Summary:
  Total requests sent:         10
  Total prompts sent:          80
  Successful requests:         10
  Failed requests:             0
  Total latency (all requests): 23658.48 ms
  Avg per request latency:     2365.73 ms
  Avg per prompt latency:      295.72 ms
  Throughput:                  3.38 prompts/second

Results using exisitng Batch Tokenization

• Launch Server Command

python -m sglang.launch_server --model-path /models/meta-llama/Llama-3.2-3B --port 30000 --host 0.0.0.0 --disable-radix-cache --disable-cuda-graph --max-prefill-tokens 131072 --chunked-prefill-size 131072 --tp 1 --dp 8 --enable-tokenizer-batch-encode

(venv) jobuser [ ~/sglang ]$ python3.10 benchmark/benchmark_batch/benchmark_batch.py 
.
.
.
Generated 480 prompts with 32000 tokens each, grouped into 60 requests of 8 prompts.

Starting benchmark: NUM_TOKENS=32000, BATCH_SIZE=8, NUM_REQUESTS=60

[Request] Sending request 1/10 with 8 prompts at 1744061155472
.
.
.
Benchmark Summary:
  Total requests sent:         60
  Total prompts sent:          480
  Successful requests:         60
  Failed requests:             0
  Total latency (all requests): 126336.63 ms
  Avg per request latency:     2105.17 ms
  Avg per prompt latency:      263.15 ms
  Throughput:                  3.80 prompts/second

• From the above benchmark we can see a good chunk of 30ms saved per prompt in a batch as measured from the client side.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[hebiao064]

Is CI down now @zhyncs

[hebiao064]

@sundar24295s qq about this:

This optimization is particularly beneficial for prefill-heavy workloads with smaller batch sizes

I feel like larger batch size will be more beneficial. Could you please explain a little bit? Thanks

[zhyncs]

Is CI down now @zhyncs

yeah just wait for @merrymercy

[sundar24295s]

@hebiao064

@sundar24295s qq about this:

This optimization is particularly beneficial for prefill-heavy workloads with smaller batch sizes

I feel like larger batch size will be more beneficial. Could you please explain a little bit? Thanks

• The current implementation tokenizes the entire batch. If the batch size is larger (say 400 or so - 400 is a rough number, it depends on the model, token length per prompt within a batch etc...), then we wait for entire batch to be tokenized before forwarding the request to GPU. During this time, GPU is idle.
• The ideal implementation would be split the big batch into smaller batches for tokenizer, tokenize them and send to the DPs which I have called out in the future work.

[hebiao064]

@hebiao064

@sundar24295s qq about this:

This optimization is particularly beneficial for prefill-heavy workloads with smaller batch sizes

I feel like larger batch size will be more beneficial. Could you please explain a little bit? Thanks

• The current implementation tokenizes the entire batch. If the batch size is larger (say 400 or so - 400 is a rough number, it depends on the model, token length per prompt within a batch etc...), then we wait for entire batch to be tokenized before forwarding the request to GPU. During this time, GPU is idle.
• The ideal implementation would be split the big batch into smaller batches for tokenizer, tokenize them and send to the DPs which I have called out in the future work.

Thanks, very clear!

[sundar24295s]

• 3 tests failed, not related to the changes in the current PR, looks flaky.
  @zhyncs Can you take a look at the PR?

[sundar24295s]

There were 3 unrelated unit test failures.

@zhyncs / @xiezhq-hermann Can you take a look at this PR?