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
+37 −2
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
ca76bfc
feat(jit): forward sub-part geometry env vars to the JIT
whn09 Aug 21, 2026
bfbdd15
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
4 changes: 4 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -337,6 +337,10 @@ The library provides some environment variables, which may be useful:
- `EP_SUPPRESS_NCCL_CHECK`: `0` or `1`, suppress NCCL version mismatch checking, `0` by default
- `EP_AVOID_RECORD_STREAM`: `0` or `1`, avoid `record_stream` on output tensors, `0` by default
- `EP_NUM_TOPK_IDX_BITS`: integer, override the number of bits for top-k index encoding, `0` (auto) by default
- `EP_NUM_SUB_PARTS`: integer, per-part sub-part count for the unordered hybrid kernel (JIT-forwarded); `0` (unset) uses the header default `2`, `1` disables sub-parts
- `EP_MIN_SUB_TOKENS`: integer, minimum tokens per sub-part before splitting (JIT-forwarded); `0` (unset) uses the header default `1`
- `EP_SM100_MIN_SUB_TOKENS`: integer, sm_100+ override for the minimum-tokens-per-sub-part gate (JIT-forwarded); `0` (unset) uses the header default `15`
- `EP_MIN_TOKENS_PER_PART`: integer, minimum tokens per scale-out part before splitting (JIT-forwarded); `0` (unset) uses the header default `15`, `1` disables the geometry clamp
- Networking
- `EP_NIC_NAME`: string, the default NIC name used to query NIC properties, `mlx5_0` by default
- `EP_OVERRIDE_RDMA_SL`: integer, override the RDMA service level index for traffic isolation
Expand Down
10 changes: 10 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,16 @@ 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
// previously required 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

@Xuan-1998 Xuan-1998 Aug 26, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should go to readme's env var list and clarify its default

@whn09 whn09 Aug 26, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in the amended head — added EP_MIN_TOKENS_PER_PART (default 15, 1 disables the clamp) to the General list.

#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

@Xuan-1998 Xuan-1998 Aug 26, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The geometry clamp is based on kNumMaxTokensPerChannel, i.e. buffer capacity, not the actual batch. If a deployment sizes one buffer for prefill and pushes decode batches through it, the clamp stays inactive and the many-tiny-parts problem remains. Sub-parts handle this at runtime via num_sub_parts_at(part_tokens). Is it possible for us to handle it at runtime too?

@whn09 whn09 Aug 26, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch — the clamp keys off kNumMaxTokensPerChannel (buffer capacity), not the per-call token count, so a caller that sizes one buffer for prefill and reuses that instantiation for decode does stay on the many-tiny-parts side. What this PR covers is the case where the callsite instantiates a kernel per shape (the bench, and callers that size one buffer per mode); for those it lands the correct geometry at JIT time.

Doing it at runtime is more invasive than num_sub_parts_at(), because kNumParts is a compile-time constant used for stack arrays (int64_t rx_part_base[kNumParts], ~L796), workspace layout (constexpr int kNumSlotsPerChannel = kNumParts * kBatchSize, L253), EP_STATIC_ASSERT(kNumParts <= 32, "One lane drains one part") at L1081, and several unrolled for (p = 0; p < kNumParts; ++p) loops. kNumSubParts can go runtime cheaply because it is a within-part ceiling and the workspace layout does not change; a runtime effective_parts would need kNumParts to stay as the allocation ceiling and every part loop / flush / signal path to key off the runtime value and skip the unused tail.

Happy to do that as a follow-up — kept out of this PR because the change surface (loops + sync + tail-skip) is much larger and warrants its own testing. Would you rather we merge this once the other comments are settled, or hold it and land runtime + compile-time together?

: ((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