[feat] support minimum token load balance in dp attention

[WANG-GH]

Motivation & Modifications

When DP attention is enabled, the system decides which DP group to dispatch a request to based on the current load of each DP group.

The load data from all DP groups is gathered once to the TP0 node, where TP0 interacts with the dispatcher via shared memory to share the load information in real time.

The load data consists of two parts:

• holding_tokens: the number of tokens currently being processed by each DP group's scheduler.

• onfly_req: the requests that have been dispatched by the dispatcher but have not yet been accepted by the scheduler (a scheduler only accepts a request when it reaches the recv_req phase).

The dispatcher sums these two parts and selects the DP group with the lowest load.

The advantage of using onfly is that we only need to track the tokens currently being processed by each worker, without worrying about whether it's MTP, overlap, etc.
I added an extra field dp_balance_id: int to the request, so we just need to gather this integer.

To enable this feature, simply add --load-balance-method minimum_tokens to the startup arguments.

Performence

Both TBT and TTFT have improvements.

python3 -m sglang.bench_serving --backend sglang --num-prompt 100

minimum_token:

python3 -m sglang.launch_server \
    --model-path /DeepSeek-R1 \
    --tp 16 \
    --nnodes 2 --node-rank 0 --trust-remote-code \
    --load-balance-method minimum_tokens \
    --dp-size 16 --cuda-graph-max-bs 128 \
    --enable-dp-attention 

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 not set   
Successful requests:                     100       
Benchmark duration (s):                  57.58     
Total input tokens:                      34308     
Total generated tokens:                  21395     
Total generated tokens (retokenized):    21292     
Request throughput (req/s):              1.74      
Input token throughput (tok/s):          595.83    
Output token throughput (tok/s):         371.57    
Total token throughput (tok/s):          967.40    
Concurrency:                             31.53     
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   18153.78  
Median E2E Latency (ms):                 15861.65  
---------------Time to First Token----------------
Mean TTFT (ms):                          4767.87   
Median TTFT (ms):                        4254.61   
P99 TTFT (ms):                           9760.97   
---------------Inter-Token Latency----------------
Mean ITL (ms):                           63.03     
Median ITL (ms):                         41.31     
P95 ITL (ms):                            57.08     
P99 ITL (ms):                            864.98    
Max ITL (ms):                            7880.85   
==================================================

rr:

python3 -m sglang.launch_server \
    --model-path /DeepSeek-R1 \
    --tp 16  \
    --nnodes 2 --node-rank 0 --trust-remote-code \
    --dp-size 16 --cuda-graph-max-bs 128 \
    --enable-dp-attention 

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 not set   
Successful requests:                     100       
Benchmark duration (s):                  63.01     
Total input tokens:                      34308     
Total generated tokens:                  21395     
Total generated tokens (retokenized):    21309     
Request throughput (req/s):              1.59      
Input token throughput (tok/s):          544.44    
Output token throughput (tok/s):         339.52    
Total token throughput (tok/s):          883.97    
Concurrency:                             36.49     
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   22992.97  
Median E2E Latency (ms):                 20915.81  
---------------Time to First Token----------------
Mean TTFT (ms):                          6945.23   
Median TTFT (ms):                        7699.87   
P99 TTFT (ms):                           14878.23  
---------------Inter-Token Latency----------------
Mean ITL (ms):                           75.57     
Median ITL (ms):                         41.34     
P95 ITL (ms):                            55.28     
P99 ITL (ms):                            1460.84   
Max ITL (ms):                            12318.38  
==================================================

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• 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.

[gemini-code-assist]

Summary of Changes

Hello @WANG-GH, 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 enhances the data parallel attention mechanism by introducing a new, more intelligent load balancing strategy. The core idea is to minimize the total token load across DP groups by considering both currently processed and in-flight requests, leveraging shared memory for efficient state synchronization. This change is expected to improve overall serving performance and resource utilization.

Highlights

• New Load Balancing Method: I've introduced a new load balancing method for Data Parallel (DP) attention called minimum_tokens. This method aims to distribute incoming requests more efficiently by dispatching them to the DP group with the lowest total token load.
• Load Calculation Logic: The 'minimum_tokens' method calculates the total load for each DP group by summing two components: holding_tokens (tokens currently being processed by the scheduler) and onfly_req (requests dispatched by the controller but not yet accepted by the scheduler). This provides a more comprehensive view of each group's real-time load.
• Shared Memory for Load Data: To facilitate real-time load balancing, I've implemented a shared memory mechanism (DPBalanceMeta). The TP0 node (dispatcher) gathers load data from all DP groups via this shared memory, allowing it to make informed dispatching decisions. A unique dp_balance_id is added to each request to track its status in the 'on-the-fly' queue.
• Performance Improvements: Initial benchmarks show that using the minimum_tokens load balancing method leads to improvements in both Time-to-First-Token (TTFT) and Total Batch Throughput (TBT) compared to the default round-robin method, indicating better overall system performance.
• Shared Memory Cleanup: I've added explicit cleanup for the POSIX shared memory segments (sglang_dp_balance_onfly_info, sglang_dp_balance_local_tokens) during process termination. This prevents shared memory leaks, which could otherwise cause issues with subsequent server launches.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a new load balancing strategy, minimum_tokens, for Data Parallel (DP) attention. This method aims to improve performance by dispatching requests to the DP group with the fewest combined holding and in-flight tokens. Key changes include:

• A new DPBalanceMeta class to manage shared memory for tracking token loads (onfly_info and local_tokens) across DP workers, including initialization and cleanup logic.
• Modifications to the DataParallelController to incorporate the new scheduling strategy, initialize DPBalanceMeta, and ensure shared memory is released.
• Updates to the Scheduler to interact with DPBalanceMeta, gather load information from DP workers, and update shared memory via the TP0 rank.
• Additions to io_struct to include dp_balance_id in request objects for tracking.
• Enhanced kill_process_tree utility to clean up the newly introduced shared memory segments.

The core logic for load balancing and shared memory management appears sound, with proper locking mechanisms. However, there's a critical issue in how local_tokens shared memory is handled (missing size header and assertion), and some areas could benefit from improved clarity with constants and variable naming. The performance benchmarks provided in the PR description show improvements with the new method.

[gemini-code-assist]

There's a critical issue in how local_tokens shared memory is handled:

1. set_shared_local_tokens writes the pickled data but does not prefix it with its size, unlike set_shared_onfly_info.
2. get_shared_local_tokens reads the entire shared memory buffer (bytes(shm.buf)) and attempts to unpickle it. If the actual pickled data is smaller than self.local_tokens_size, pickle.loads will try to interpret trailing garbage data, leading to potential errors or incorrect deserialization.
3. set_shared_local_tokens does not assert that the serialized_data fits within self.local_tokens_size.

To fix this, set_shared_local_tokens should be modified to first write the size of the serialized data (using struct.pack like in set_shared_onfly_info), and get_shared_local_tokens should read this size to correctly slice the buffer before unpickling. An assertion for data size should also be added to set_shared_local_tokens.

[gemini-code-assist]

The calculation for onfly_info_size and local_tokens_size uses a magic number 8 (commented as dByte).

1. For onfly_info_size, 512 represents max_onfly_req_per_worker. The * 8 seems to be an estimate for the pickled size of dictionary entries. While the assertion in set_shared_onfly_info provides safety, consider making MAX_ONFLY_REQ_PER_WORKER = 512 and ESTIMATED_BYTES_PER_ENTRY = 8 (or a more descriptive name if 8 has a specific meaning) as named constants for better readability and maintainability.

2. For local_tokens_size = num_workers * 8, this calculation might be too small for storing a pickled list of num_workers integers, especially considering Python's pickle overhead. For example, pickle.dumps([0]*16) is 50 bytes, not 16*8=128. More importantly, set_shared_local_tokens currently lacks an assertion to check if the pickled data fits within local_tokens_size. This, combined with the issue in get_shared_local_tokens (see separate comment), is problematic.

[gemini-code-assist]

This log message uses logger.info and includes potentially large data structures (local_tokens, onfly_info). If requests are frequent, this could lead to excessive logging and performance overhead. Consider changing this to logger.debug or making it conditional, for example, logging only every N requests or if a specific debug flag is enabled.

[gemini-code-assist]

This line self.shm1.name accesses the name attribute of the shared memory object but does not use the value. It has no side effect and can be removed if it's not intended for a specific purpose (e.g., debugging during development).

[gemini-code-assist]

The parameter holding_tokens_list in gather_dp_balance_info is misleading. Based on its usage in handle_dp_balance_data (line 1854-1855), it's a single integer representing the current worker's holding tokens, not a list.

Later, within gather_dp_balance_info (line 1814), holding_tokens_list is re-assigned to be a list of holding tokens gathered from all workers if self.tp_rank == 0.

This dual meaning and misnaming can cause confusion. Consider renaming the parameter to something like current_worker_holding_tokens: int and the locally gathered list to all_workers_holding_tokens: List[int].

[gemini-code-assist]

The values 511 and 512 (for gather_tensor_size) are used here and seem related to max_onfly_req_per_worker defined implicitly in DPBalanceMeta. It would be better to define these as named constants, possibly in data_parallel_meta.py or a shared constants module, and import them here. This improves readability and makes it easier to update if the underlying limits change.

[gemini-code-assist]

The return type hint Union[None, List[List[int]]] for gather_dp_balance_info is incorrect. The function actually returns a tuple: (gathered_id_list_per_worker, holding_tokens_list).
Based on the logic:

• If self.tp_rank != 0, it returns (None, input_argument_holding_tokens).
• If self.tp_rank == 0, it returns (List[List[int]], List[int]) (where the second list is the gathered holding tokens for all workers).

A more accurate type hint would be Tuple[Optional[List[List[int]]], List[int]].

Suggested change

	def gather_dp_balance_info(holding_tokens_list) -> Union[None, List[List[int]]]:
	"""gather recv_dp_balance_id_this_term and holding tokens per worker for dp balance"""
	recv_list = self.recv_dp_balance_id_this_term
	assert len(recv_list) <= 511, (
	"The number of requests received this round is too large. "
	"Please increase gather_tensor_size and onfly_info_size."
	)
	gather_tensor_size = 512
	# recv_tensor: | holding_tokens | len(recv_dp_balance_id) | recv_dp_balance_ids
	recv_tensor = torch.zeros(gather_tensor_size, dtype=torch.int32)
	recv_tensor[0] = holding_tokens_list
	recv_tensor[1] = len(
	recv_list
	) # The first element is the length of the list.
	recv_tensor[2 : len(recv_list) + 2] = torch.tensor(
	recv_list, dtype=torch.int32
	)
	if self.tp_rank == 0:
	gathered_list = [
	torch.zeros(gather_tensor_size, dtype=torch.int32)
	for _ in range(self.balance_meta.num_workers)
	]
	else:
	gathered_list = None
	torch.distributed.gather(
	recv_tensor, gathered_list, group=self.tp_cpu_group
	)
	gathered_id_list_per_worker = None
	if self.tp_rank == 0:
	gathered_id_list_per_worker = []
	holding_tokens_list = []
	for tensor in gathered_list:
	holding_tokens_list.append(tensor[0].item())
	list_length = tensor[1].item()
	gathered_id_list_per_worker.append(
	tensor[2 : list_length + 2].tolist()
	)
	return gathered_id_list_per_worker, holding_tokens_list
	def gather_dp_balance_info(current_worker_holding_tokens: int) -> Tuple[Optional[List[List[int]]], List[int]]:

[Edenzzzz]

might be more formal to name these as "on_the_fly" or "in_flight"?

[WANG-GH]

Sure, we can first discuss the current implementation of the algorithm, and then I’ll make a unified update accordingly.

[Edenzzzz]

I see, previously it only uses round robin without checking token count

[JustinTong0323]

I think you should clean the commit history🤣

[WANG-GH]

I've just cleaned up the commit history and resolved the conflicts with the main branch. Could you please review this PR?
@JustinTong0323

[JustinTong0323]

Plz fix the lint, you could ref to contribution guide.

[WANG-GH]

Hi, i finished the lint error, could you please trigger the ci?

[Qiaolin-Yu]

Could you add explanation for this value?

[WANG-GH]

it means the maximum size of the tensor used for gathering data.

[Qiaolin-Yu]

Could you explain the meaning of this variable?

[WANG-GH]

sure, this variable corresponds to the balance_id in TokenizedGenerateReqInput.
We use it to to control the number of onfly tokens (requests dispatched to workers but not yet received).

[Qiaolin-Yu]

I was thinking if we could abstract some methods to python/sglang/srt/utils.py and python/sglang/srt/distributed/parallel_state.py ?

[WANG-GH]

sure, i moved it to the manager/utils.py

[WANG-GH]

I tested on 2*node H20 with 4096 reqs

# test command
python3 -m sglang.bench_serving --backend sglang --num-prompts 4096 --dataset-path /sgl-workspace/sharegpt.json

rr：

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 not set   
Successful requests:                     4096      
Benchmark duration (s):                  263.31    
Total input tokens:                      1294121   
Total generated tokens:                  787217    
Total generated tokens (retokenized):    783897    
Request throughput (req/s):              15.56     
Input token throughput (tok/s):          4914.74   
Output token throughput (tok/s):         2989.65   
Total token throughput (tok/s):          7904.38   
Concurrency:                             2123.68   
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   136522.09 
Median E2E Latency (ms):                 130958.79 
---------------Time to First Token----------------
Mean TTFT (ms):                          60330.23  
Median TTFT (ms):                        63801.38  
P99 TTFT (ms):                           113752.40 
---------------Inter-Token Latency----------------
Mean ITL (ms):                           396.08    
Median ITL (ms):                         143.72    
P95 ITL (ms):                            654.97    
P99 ITL (ms):                            1512.92   
Max ITL (ms):                            105775.01 
==================================================

minimum token：

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 not set   
Successful requests:                     4096      
Benchmark duration (s):                  262.36    
Total input tokens:                      1294121   
Total generated tokens:                  787217    
Total generated tokens (retokenized):    783888    
Request throughput (req/s):              15.61     
Input token throughput (tok/s):          4932.61   
Output token throughput (tok/s):         3000.52   
Total token throughput (tok/s):          7933.12   
Concurrency:                             2136.54   
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   136851.55 
Median E2E Latency (ms):                 129838.93 
---------------Time to First Token----------------
Mean TTFT (ms):                          61737.04  
Median TTFT (ms):                        64347.30  
P99 TTFT (ms):                           111977.50 
---------------Inter-Token Latency----------------
Mean ITL (ms):                           390.17    
Median ITL (ms):                         143.37    
P95 ITL (ms):                            567.51    
P99 ITL (ms):                            1079.40   
Max ITL (ms):                            104110.49 
==================================================

[ltdo111]

I've been following your PR for a long time. I was wondering when it will be merged into the main repository. Do you have a plan for that?

[WANG-GH]

I've been following your PR for a long time. I was wondering when it will be merged into the main repository. Do you have a plan for that?

Thank you for the reply! I also hope to merge it into the main branch as soon as possible. Over the past few weeks, I've been actively communicating with Qiaolin and Cheng Wan from the SGLang community, and they have provided many great suggestions for this PR. We are currently confirming with the PD separation team to check if there are any conflicts between our designs.
@ch-wan @Qiaolin-Yu

[ltdo111]

I've been following your PR for a long time. I was wondering when it will be merged into the main repository. Do you have a plan for that?

Thank you for the reply! I also hope to merge it into the main branch as soon as possible. Over the past few weeks, I've been actively communicating with Qiaolin and Cheng Wan from the SGLang community, and they have provided many great suggestions for this PR. We are currently confirming with the PD separation team to check if there are any conflicts between our designs. @ch-wan @Qiaolin-Yu

I've read through your PR thoroughly and have two points I'd like to confirm with you:

The first question: There is a get_load method in the scheduler, but it isn't reused in this PR. After reading through your code, my understanding is that the current minimum token load balancing method takes into account the load brought by output token IDs during the decode phase, making the load statistics more accurate. I wonder if this understanding is correct.

The second question is about the implementation of data_parallel_controller and IPC communication. The community originally adopted a unified communication method using zmq, but this PR introduces a new mp API to implement communication between dpc and the scheduler. Have you considered switching the communication method from mp to IPC communication? This would be more consistent with the community style, provide a single unified implementation for inter-process communication, and make it easier for subsequent open-source contributors to read the code.

[PanXun2]

@WANG-GH Glad to know this is actively developed. Currently, the Router has the load balancer function which get_loads from each sglang instance. It has the function for collecting workloads. Do you have any plan to reuse such loads information? Or any efforts to make code better reused for these similar functions?

[WANG-GH]

I've read through your PR thoroughly and have two points I'd like to confirm with you:

The first question: There is a get_load method in the scheduler, but it isn't reused in this PR. After reading through your code, my understanding is that the current minimum token load balancing method takes into account the load brought by output token IDs during the decode phase, making the load statistics more accurate. I wonder if this understanding is correct.

The second question is about the implementation of data_parallel_controller and IPC communication. The community originally adopted a unified communication method using zmq, but this PR introduces a new mp API to implement communication between dpc and the scheduler. Have you considered switching the communication method from mp to IPC communication? This would be more consistent with the community style, provide a single unified implementation for inter-process communication, and make it easier for subsequent open-source contributors to read the code.

Thanks for ask.

For the first question: I hadn't noticed the get_load function before. When calculating the load, this function, combined with the current batch size, can better accommodate PD separation.

For the second question: IPC communication is more suitable for the producer-consumer model, but the current controller needs to frequently check the global instance's status. This programming paradigm is more like using locks to maintain a global state, which doesn't align well with the producer-consumer model. Changing to IPC would be quite awkward.

[WANG-GH]

@WANG-GH Glad to know this is actively developed. Currently, the Router has the load balancer function which get_loads from each sglang instance. It has the function for collecting workloads. Do you have any plan to reuse such loads information? Or any efforts to make code better reused for these similar functions?

I hadn't noticed the get_load function before. When calculating the load, this function, combined with the current batch size, can better accommodate PD separation. I will refactor the handle_dp_balance_data func to use get_loads before gather.

[ltdo111]

I've read through your PR thoroughly and have two points I'd like to confirm with you:
The first question: There is a get_load method in the scheduler, but it isn't reused in this PR. After reading through your code, my understanding is that the current minimum token load balancing method takes into account the load brought by output token IDs during the decode phase, making the load statistics more accurate. I wonder if this understanding is correct.
The second question is about the implementation of data_parallel_controller and IPC communication. The community originally adopted a unified communication method using zmq, but this PR introduces a new mp API to implement communication between dpc and the scheduler. Have you considered switching the communication method from mp to IPC communication? This would be more consistent with the community style, provide a single unified implementation for inter-process communication, and make it easier for subsequent open-source contributors to read the code.

Thanks for ask.

For the first question: I hadn't noticed the get_load function before. When calculating the load, this function, combined with the current batch size, can better accommodate PD separation.

For the second question: IPC communication is more suitable for the producer-consumer model, but the current controller needs to frequently check the global instance's status. This programming paradigm is more like using locks to maintain a global state, which doesn't align well with the producer-consumer model. Changing to IPC would be quite awkward.

thx for reply, i see。

[ollybbmonster]

Thx for your work on this feature, it helps me a lot.

[ollybbmonster]

I’m wondering how long onfly_reqs can actually survive.
It seems the scheduler receives reqs from DPC almost immediately.
Meanwhile, onfly_reqs are appended in process_input_requests then excluded at the end of get_next_batch_to_run per 40 iterations.

Given this flow, the comment here might be misleading.

[ollybbmonster]

Not sure if this is intentional, but in both set_* functions, passing a Python list directly into a multiprocessing.Manager().List() could replace the managed object, losing the cross-process synchronization.

[ollybbmonster]

Should this be assert len(recv_list) < 511?
Or recv_tensor could be 1 + 1 + 511 in length, which would exceed 512.

Also, as the Gemini bot mentioned, holding_tokens_list is misleading when it’s actually length of tokens rather than a list — especially since it’s later mixed with real list operations.

[whybeyoung]

LGTM