[PR #599] [MERGED] Support topk_weights, recv rank info, and overlap in low-latency dispatch #652

New issue

Closed

opened 2026-10-27 21:04:54 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-27 21:04:54 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/599
Author: @nvcastet
Created: 4/1/2026
Status: ✅ Merged
Merged: 4/14/2026
Merged by: @jershi425

Base: hybrid-ep ← Head: add_ll_topk_weight_v2

📝 Commits (1)

fc69fb0 Support topk_weights, recv rank info, and overlap in ll dispatch

📊 Changes

6 files changed (+159 additions, -52 deletions)

View changed files

📝 csrc/deep_ep.cpp (+29 -19)
📝 csrc/deep_ep.hpp (+3 -2)
📝 csrc/kernels/api.cuh (+4 -1)
📝 csrc/kernels/internode_ll.cu (+73 -19)
📝 csrc/kernels/launch.cuh (+22 -1)
📝 deep_ep/buffer.py (+28 -10)

📄 Description

Summary

Enables https://github.com/sgl-project/sglang/pull/21877.

Adds optional topk_weights parameter to low_latency_dispatch, allowing topk weights to be sent alongside token data in the P2P message header (packed as int4{token_idx, topk_weight_bits, 0, 0})
Returns two new output tensors when topk_weights is provided: packed_recv_topk_weights (float32) and packed_recv_rank_info (int32), giving the receiver both the weight and source rank for each token
Backward compatible: returns 5 elements when topk_weights=None, 7 elements when provided

Limitations

TopK weights dispatch is only supported on the P2P (NVLink) path; an assertion fires if the RDMA (non-P2P / IBGDA) fallback is hit

Files changed

csrc/kernels/internode_ll.cu — new kUseTopKWeights template param, extended header write, recv-side unpacking
csrc/kernels/launch.cuh — SETUP_OVERLAP_LAUNCH_CONFIG macro
csrc/kernels/api.cuh — updated function signatures
csrc/deep_ep.cpp — plumbing for new args/outputs
csrc/deep_ep.hpp — updated return type and signature
deep_ep/buffer.py — Python API: new topk_weights param, conditional return format, updated docstring

Test plan

Run test_low_latency.py with and without topk_weights on a multi-GPU NVLink node
Verify backward compatibility: existing callers without topk_weights get the same 5-element return
Validate received topk_weights and rank_info match what was sent

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/599 **Author:** [@nvcastet](https://github.com/nvcastet) **Created:** 4/1/2026 **Status:** ✅ Merged **Merged:** 4/14/2026 **Merged by:** [@jershi425](https://github.com/jershi425) **Base:** `hybrid-ep` ← **Head:** `add_ll_topk_weight_v2` --- ### 📝 Commits (1) - [`fc69fb0`](https://github.com/deepseek-ai/DeepEP/commit/fc69fb0d2910c41b4c0cb956920654df52105aff) Support topk_weights, recv rank info, and overlap in ll dispatch ### 📊 Changes **6 files changed** (+159 additions, -52 deletions) <details> <summary>View changed files</summary> 📝 `csrc/deep_ep.cpp` (+29 -19) 📝 `csrc/deep_ep.hpp` (+3 -2) 📝 `csrc/kernels/api.cuh` (+4 -1) 📝 `csrc/kernels/internode_ll.cu` (+73 -19) 📝 `csrc/kernels/launch.cuh` (+22 -1) 📝 `deep_ep/buffer.py` (+28 -10) </details> ### 📄 Description ## Summary Enables https://github.com/sgl-project/sglang/pull/21877. - Adds optional `topk_weights` parameter to `low_latency_dispatch`, allowing topk weights to be sent alongside token data in the P2P message header (packed as `int4{token_idx, topk_weight_bits, 0, 0}`) - Returns two new output tensors when `topk_weights` is provided: `packed_recv_topk_weights` (float32) and `packed_recv_rank_info` (int32), giving the receiver both the weight and source rank for each token - Backward compatible: returns 5 elements when `topk_weights=None`, 7 elements when provided ## Limitations - TopK weights dispatch is only supported on the P2P (NVLink) path; an assertion fires if the RDMA (non-P2P / IBGDA) fallback is hit ## Files changed - `csrc/kernels/internode_ll.cu` — new `kUseTopKWeights` template param, extended header write, recv-side unpacking - `csrc/kernels/launch.cuh` — `SETUP_OVERLAP_LAUNCH_CONFIG` macro - `csrc/kernels/api.cuh` — updated function signatures - `csrc/deep_ep.cpp` — plumbing for new args/outputs - `csrc/deep_ep.hpp` — updated return type and signature - `deep_ep/buffer.py` — Python API: new `topk_weights` param, conditional return format, updated docstring ## Test plan - [ ] Run `test_low_latency.py` with and without `topk_weights` on a multi-GPU NVLink node - [ ] Verify backward compatibility: existing callers without `topk_weights` get the same 5-element return - [ ] Validate received `topk_weights` and `rank_info` match what was sent --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-27 21:04:54 +00:00

closed this issue
added the
pull-request
label

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#652

No description provided.

Rows
Columns