Skip to content

perf(hybrid): don't split a channel into more parts than it has tokens for #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
whn09 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

perf(hybrid): don't split a channel into more parts than it has tokens for #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5118c2e
feat(jit): forward sub-part geometry env vars to the JIT
whn09 Aug 21, 2026
b097b03
perf(hybrid): don't split a channel into more parts than it has token…
whn09 Aug 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions csrc/jit/compiler.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,17 @@ class Compiler {
// TODO: make it more general, e.g. `EP_JIT_EXTRA_FLAGS`
if (int num_topk_idx_bits = get_env("EP_NUM_TOPK_IDX_BITS", 0); num_topk_idx_bits != 0)
flags += fmt::format(" -DEP_NUM_TOPK_IDX_BITS={}", num_topk_idx_bits);

// Part / sub-part geometry defaults in `hybrid_dispatch_unordered.cuh`. They are
// device-only (no
// host caller reads them), so forwarding them as JIT defines cannot desync host and device,
// and `flags` is part of `kernel_signature` below, so a change re-JITs rather than
// reusing a stale cubin. Tuning them per network/arch currently requires editing the
// header and reinstalling.
for (const auto& name: {"EP_NUM_SUB_PARTS", "EP_MIN_SUB_TOKENS", "EP_SM100_MIN_SUB_TOKENS",
"EP_MIN_TOKENS_PER_PART"})
if (int v = get_env(name, 0); v != 0)
flags += fmt::format(" -D{}={}", name, v);
}

virtual ~Compiler() = default;
Expand Down
25 changes: 23 additions & 2 deletions deep_ep/include/deep_ep/impls/hybrid_dispatch_unordered.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,20 @@ static constexpr int kMinSubTokensDefault = (EP_MIN_SUB_TOKENS) > 1 ? (EP_MIN_SU
#define EP_SM100_MIN_SUB_TOKENS 15
#endif

// Minimum tokens a scale-out part must carry to be worth its own `flush_part` put, i.e. the
// part-level analogue of `kMinSubTokensDefault` above. `compute_part_allocation()` only ever
// caps the part count from ABOVE when the indexed-signal budget is tight, and that budget is
// loosest exactly when a channel holds the fewest tokens, so small-batch shapes settle on
// `kMaxParts` -- the worst end of the axis. Set to 1 to disable the clamp entirely, i.e. to
// restore the previous behaviour exactly (a value of 1 must SHORT-CIRCUIT rather than divide by
// one: `kNumMaxTokensPerChannel / 1` still clamps whenever a channel holds fewer tokens than the
// budget allows parts, which is a different geometry from the old code and not a control).
#ifndef EP_MIN_TOKENS_PER_PART
#define EP_MIN_TOKENS_PER_PART 15
#endif

static constexpr int kMinTokensPerPart = (EP_MIN_TOKENS_PER_PART) > 1 ? (EP_MIN_TOKENS_PER_PART) : 1;

template <int kNumSubParts, int kMinSubTokens = kMinSubTokensDefault>
__device__ __host__ __forceinline__ int num_sub_parts_at(const int& part_tokens) {
if constexpr (kNumSubParts <= 1) {
Expand Down Expand Up @@ -140,9 +154,16 @@ template <bool kDoCPUSync,
int kNumScaleupRanksPerLane = math::constexpr_ceil_div(kNumScaleupRanks, 32),
int kNumChannelsPerSM = kNumScaleoutWarps,
int kNumChannels = kNumScaleoutWarps * kNumSMs,
int kNumParts = gin_alloc::constexpr_num_parts(
kNumGinSignals, kNumSMs, kNumQPs, (kNumNotifyWarps > 0), kNumScaleoutWarps),
int kNumMaxTokensPerChannel = math::constexpr_ceil_div(kNumMaxTokensPerRank, kNumChannels),
int kNumBudgetParts = gin_alloc::constexpr_num_parts(
kNumGinSignals, kNumSMs, kNumQPs, (kNumNotifyWarps > 0), kNumScaleoutWarps),
// NOTES: the parentheses around the comparison are load-bearing -- an unparenthesized
// `>` inside a template parameter list closes the list instead of comparing (same
// reason `(kNumNotifyWarps > 0)` above is wrapped)
int kNumGeomParts = kMinTokensPerPart <= 1 ? kNumBudgetParts
: ((kNumMaxTokensPerChannel / kMinTokensPerPart > 1)
? kNumMaxTokensPerChannel / kMinTokensPerPart : 1),
int kNumParts = kNumBudgetParts < kNumGeomParts ? kNumBudgetParts : kNumGeomParts,
int kPartSize = math::constexpr_ceil_div(kNumMaxTokensPerChannel, kNumParts),
int kBatchSize = kPartSize,
int kNumSubParts = kNumSubPartsDefault < kBatchSize ? kNumSubPartsDefault : kBatchSize,
Expand Down