deepseek-ai/DeepEP

Fork 0

mirror of https://github.com/deepseek-ai/DeepEP.git synced 2026-11-05 00:47:32 +00:00

[GH-ISSUE #241] Maybe a Method to Reduce Memory Consumption in Low-Latency Mode #186

New issue

Closed

opened 2026-10-26 09:56:03 +00:00 by bsakaguchi · 2 comments

bsakaguchi commented

2026-10-26 09:56:03 +00:00

Owner

Originally created by @yulangz on GitHub (Jun 21, 2025).
Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/241

In low-latency mode, since it's impossible to predict the expert selected for each token beforehand, buffer allocation must be based on the worst-case scenario – assuming all tokens will be sent to a given expert. In reality, a significant portion of the pre-allocated buffer remains unused. I have some ideas that might reduce video memory consumption in low-latency mode:

During the dispatch phase, reserve a dispatch_recv_data_buffer_bytes size of redund_ratio * num_topk * num_local_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_dispatch_msg, instead of num_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_dispatch_msg. Here, redund_ratio is an empirical value representing the potential extra selections an expert might receive under relatively imbalanced system loads. Assuming redund_ratio=2, num_topk=8, num_experts=256, num_local_experts=1, this would reduce memory usage by a factor of 16.

In the actual dispatch phase, to ensure sufficient memory space, the following approach can be adopted:

Each expert maintains a recvBuffer, along with two control fields: recv_token_num and finish_rank_num. The token capacity of recv_buffer len_{buffer} is num_max_dispatch_tokens_per_rank * num_topk * redund_ratio.
On each rank, first count the number of tokens num_{e_i} to be sent to each expert e_i.
Atomically add num_{e_i} to the target expert's recv_token_num field, and get the return value slot_{end}.
If slot_{end} < len_{buffer}, it indicates sufficient space in the target expert's recv_buffer. Copy the tokens to the location recvBuffer[slot_{end} - num_{e_i} : slot_{end}].
Otherwise, it indicates insufficient space in the target expert's recvBuffer. The overflow amount is slot_{end} - len_{buffer}, meaning these tokens cannot be sent to the target expert. We need to:
1. Calculate num_{block} = min(num_{e_i}, slot_{end} - len_{buffer}), which is the number of tokens that cannot be sent to the target expert.
2. Select num_{block} tokens and mark them as failed sends.
3. For the remaining tokens, send them to the target expert's recvBuffer[slot_{end} - num_{e_i} : len_{buffer}].
Regardless of the case, after sending, atomically increment the target expert's finish_rank_num by 1. In the LOW_LATENCY_DISPATCH_RECV phase, wait for finish_rank_num == num_ranks to indicate reception completion.

Pseudocode:

// Buffer space per expert
num_slot_num_per_expert = redund_ratio * num_topk * num_max_dispatch_tokens_per_rank;
dispatch_recv_data_buffer_bytes = num_slot_num_per_expert * num_local_experts * num_bytes_per_dispatch_msg;
T dispatch_recv_buffer[dispatch_recv_data_buffer_bytes];
int token_num_per_rank[num_ranks];
int recv_token_num = 0;
int finish_rank_num = 0;

__global__ void low_latency_dispatch(...) {
    // Count total token number to be sent to expert e_i
    int num_i = count_total_token_num(e_i); 
    // Atomically add num_i to recv_token_num of target expert to allocate slot
    int slot_end = rdma_atomic_add(e_i, recv_token_num, num_i);
    int real_copy_num;
    if (slot_end < num_slot_num_per_expert) {
        // Sufficient slot on target expert: copy all tokens
        rdma_copy_to(e_i, dispatch_recv_buffer[slot_end - num_i: slot_end], token[:]);
        real_copy_num = num_i;
    } else {
        // Insufficient slot: discard some tokens, copy the rest
        int discard_token_num = min(num_i, slot_end - num_slot_num_per_expert);
        mask_discard(token[num_i - discard_token_num: num_i]);
        real_copy_num = num_i - discard_token_num;
        rdma_copy_to(e_i, dispatch_recv_buffer[num_slot_num_per_expert - real_copy_num: num_slot_num_per_expert], token[:real_copy_num]);
    }
    // Notify target expert of completion
    rdma_write(e_i, token_num_per_rank[this_rank], real_copy_num);
    rdma_atomic_add(e_i, finish_rank_num, 1);
}

How to handle discarded tokens? I envision two approaches:

When using micro-batching and currently at batch0, have batch1 handle the discarded tokens.
Alternatively, inform the inference system about the discarded tokens. The system skips generating output for them in this round and processes them in the next. However, this requires additional support from the inference system.
In systems with relatively balanced expert loads, token discarding should occur infrequently, minimizing its overall performance impact.

By reducing dispatch_recv_data_buffer_bytes, the sizes of combine_send_buffer_bytes and packed_recv_x (https://github.com/deepseek-ai/DeepEP/blob/7b0c25f864cc7a2d95d7e3c3c2b7c13b22116063/csrc/deep_ep.cpp#L1089C33-L1089C38) can also be reduced.

Assuming num_max_dispatch_tokens_per_rank=128, hidden=7168, num_experts=256, num_ranks=1, redund_ratio=2, this method could reduce the buffer size for LowLatencyLayout from 1.75GB down to 112MB.

By the way, although this method reduces memory consumption, it introduces system complexity. Furthermore, the original 1.75GB also seems like an acceptable size.

I look forward to discussing this with you!

Originally created by @yulangz on GitHub (Jun 21, 2025). Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/241 In low-latency mode, since it's impossible to predict the expert selected for each token beforehand, buffer allocation must be based on the worst-case scenario – assuming all tokens will be sent to a given expert. In reality, a significant portion of the pre-allocated buffer remains unused. I have some ideas that might reduce video memory consumption in low-latency mode: During the dispatch phase, reserve a `dispatch_recv_data_buffer_bytes` size of `redund_ratio * num_topk * num_local_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_dispatch_msg`, instead of `num_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_dispatch_msg`. Here, `redund_ratio` is an empirical value representing the potential extra selections an expert might receive under relatively imbalanced system loads. Assuming `redund_ratio=2, num_topk=8, num_experts=256, num_local_experts=1`, this would reduce memory usage by a factor of 16. In the actual dispatch phase, to ensure sufficient memory space, the following approach can be adopted: 1. Each expert maintains a recvBuffer, along with two control fields: `recv_token_num` and `finish_rank_num`. The token capacity of recv_buffer $len_{buffer}$ is `num_max_dispatch_tokens_per_rank * num_topk * redund_ratio`. 2. On each rank, first count the number of tokens $num_{e_i}$ to be sent to each expert $e_i$. 3. Atomically add $num_{e_i}$ to the target expert's `recv_token_num` field, and get the return value $slot_{end}$. 4. If $slot_{end} < len_{buffer}$, it indicates sufficient space in the target expert's recv_buffer. Copy the tokens to the location $recvBuffer[slot_{end} - num_{e_i} : slot_{end}]$. 5. Otherwise, it indicates insufficient space in the target expert's recvBuffer. The overflow amount is $slot_{end} - len_{buffer}$, meaning these tokens cannot be sent to the target expert. We need to: 1. Calculate $num_{block} = min(num_{e_i}, slot_{end} - len_{buffer})$, which is the number of tokens that cannot be sent to the target expert. 2. Select $num_{block}$ tokens and mark them as failed sends. 3. For the remaining tokens, send them to the target expert's $recvBuffer[slot_{end} - num_{e_i} : len_{buffer}]$. 6. Regardless of the case, after sending, atomically increment the target expert's `finish_rank_num` by 1. In the `LOW_LATENCY_DISPATCH_RECV` phase, wait for `finish_rank_num == num_ranks` to indicate reception completion. Pseudocode: ```c++ // Buffer space per expert num_slot_num_per_expert = redund_ratio * num_topk * num_max_dispatch_tokens_per_rank; dispatch_recv_data_buffer_bytes = num_slot_num_per_expert * num_local_experts * num_bytes_per_dispatch_msg; T dispatch_recv_buffer[dispatch_recv_data_buffer_bytes]; int token_num_per_rank[num_ranks]; int recv_token_num = 0; int finish_rank_num = 0; __global__ void low_latency_dispatch(...) { // Count total token number to be sent to expert e_i int num_i = count_total_token_num(e_i); // Atomically add num_i to recv_token_num of target expert to allocate slot int slot_end = rdma_atomic_add(e_i, recv_token_num, num_i); int real_copy_num; if (slot_end < num_slot_num_per_expert) { // Sufficient slot on target expert: copy all tokens rdma_copy_to(e_i, dispatch_recv_buffer[slot_end - num_i: slot_end], token[:]); real_copy_num = num_i; } else { // Insufficient slot: discard some tokens, copy the rest int discard_token_num = min(num_i, slot_end - num_slot_num_per_expert); mask_discard(token[num_i - discard_token_num: num_i]); real_copy_num = num_i - discard_token_num; rdma_copy_to(e_i, dispatch_recv_buffer[num_slot_num_per_expert - real_copy_num: num_slot_num_per_expert], token[:real_copy_num]); } // Notify target expert of completion rdma_write(e_i, token_num_per_rank[this_rank], real_copy_num); rdma_atomic_add(e_i, finish_rank_num, 1); } ``` How to handle discarded tokens? I envision two approaches: 1. When using micro-batching and currently at batch0, have batch1 handle the discarded tokens. 2. Alternatively, inform the inference system about the discarded tokens. The system skips generating output for them in this round and processes them in the next. However, this requires additional support from the inference system. In systems with relatively balanced expert loads, token discarding should occur infrequently, minimizing its overall performance impact. By reducing `dispatch_recv_data_buffer_bytes`, the sizes of `combine_send_buffer_bytes` and packed_recv_x ([https://github.com/deepseek-ai/DeepEP/blob/7b0c25f864cc7a2d95d7e3c3c2b7c13b22116063/csrc/deep_ep.cpp#L1089C33-L1089C38](https://github.com/deepseek-ai/DeepEP/blob/7b0c25f864cc7a2d95d7e3c3c2b7c13b22116063/csrc/deep_ep.cpp#L1089C33-L1089C38)) can also be reduced. Assuming num_max_dispatch_tokens_per_rank=128, hidden=7168, num_experts=256, num_ranks=1, redund_ratio=2, this method could reduce the buffer size for `LowLatencyLayout` from 1.75GB down to 112MB. By the way, although this method reduces memory consumption, it introduces system complexity. Furthermore, the original 1.75GB also seems like an acceptable size. I look forward to discussing this with you!

bsakaguchi closed this issue

2026-10-26 09:56:03 +00:00

bsakaguchi commented

2026-10-26 09:56:04 +00:00

Author

Owner

@LyricZhao commented on GitHub (Jun 27, 2025):

I think it is a great idea! For real online services, the buffer size of low-latency kernels is not a problem. But for normal kernels, using a communication queue will lead to performance cost, using an estimated buffer size is a great idea👍🏻

We plan to do a full refactor later in this year, we must have such ideas to be adopted.

@LyricZhao commented on GitHub (Jun 27, 2025): I think it is a great idea! For real online services, the buffer size of low-latency kernels is not a problem. But for normal kernels, using a communication queue will lead to performance cost, using an estimated buffer size is a great idea👍🏻 We plan to do a full refactor later in this year, we must have such ideas to be adopted.

bsakaguchi commented

2026-10-26 09:56:04 +00:00

Author

Owner

@yulangz commented on GitHub (Jul 4, 2025):

Thank you for your appreciation! I am looking forward to seeing refactoring in the future!

@yulangz commented on GitHub (Jul 4, 2025): Thank you for your appreciation! I am looking forward to seeing refactoring in the future!

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

deepseek-ai/DeepEP#186

No description provided.

Rows
Columns