From ca76bfce2331d36345c1fb0c00a0f2e273bb6b0a Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 21 Aug 2026 05:36:10 +0000 Subject: [PATCH 1/2] feat(jit): forward sub-part geometry env vars to the JIT `hybrid_dispatch_unordered.cuh` gates the sub-part geometry behind `#ifndef` (`EP_NUM_SUB_PARTS` 2, `EP_MIN_SUB_TOKENS` 1, `EP_SM100_MIN_SUB_TOKENS` 15), but nothing in the tree sets those macros, so the only way to try a different split is to edit the header and reinstall. Forward the three names as JIT `-D` flags, following the `EP_NUM_TOPK_IDX_BITS` block immediately above (and its `EP_JIT_EXTRA_FLAGS` TODO). All three are device-only -- no host translation unit reads them -- so a JIT-only define cannot desync host and device sizing. `flags` is part of `kernel_signature`, so changing the env re-JITs instead of serving a cached cubin. Unset => no behaviour change. (cherry picked from commit 5118c2e3f97920841567eafd868b850a091ff3a0) --- README.md | 3 +++ csrc/jit/compiler.hpp | 9 +++++++++ 2 files changed, 12 insertions(+) diff --git a/README.md b/README.md index 078a15146..dadda17f4 100644 --- a/README.md +++ b/README.md @@ -337,6 +337,9 @@ 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` - 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 diff --git a/csrc/jit/compiler.hpp b/csrc/jit/compiler.hpp index ad01b3cac..d471e736e 100644 --- a/csrc/jit/compiler.hpp +++ b/csrc/jit/compiler.hpp @@ -70,6 +70,15 @@ 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); + + // 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"}) + if (int v = get_env(name, 0); v != 0) + flags += fmt::format(" -D{}={}", name, v); } virtual ~Compiler() = default; From bfbdd15ff448783f877cb2210cb3246c8452b05e Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 21 Aug 2026 05:36:10 +0000 Subject: [PATCH 2/2] perf(hybrid): don't split a channel into more parts than it has tokens for `kNumParts` -- how many `flush_part` puts a channel's tokens leave in -- is chosen today by `compute_part_allocation()` alone, which only ever caps the count from ABOVE when the GIN indexed-signal budget is tight. It is never lowered because the geometry asks for it: `kNumParts` is never compared against `kNumMaxTokensPerChannel`, and there is no minimum-tokens-per-part threshold. The budget is loosest exactly when a channel holds the fewest tokens (low `--num-sms`, small batch), so decode shapes land on `kMaxParts` -- the worst end of the axis -- with no way to opt out. At 128 tokens / 12 SMs a channel holds 3 tokens and is described as 4 parts x 1 token: the last part is always empty, and 3 tokens leave as three separate single-token puts instead of one 3-token put. Give parts the guard sub-parts already have. Sub-parts have both a clamp of `kNumSubParts` to `kBatchSize` and `EP_SM100_MIN_SUB_TOKENS` refusing to sub-split a part too small to be worth it; parts have neither. `kMinTokensPerPart` defaults to 15 (copied from the sub-token precedent in the same file) and is overridable by `EP_MIN_TOKENS_PER_PART`. `EP_MIN_TOKENS_PER_PART=1` short-circuits to the old value, so it is an exact in-image control rather than an approximation. `kNumMaxTokensPerChannel` moves above the part count in the template list; it depends only on already-declared parameters. (cherry picked from commit b097b03799533c911a1a594fdeb82375fa8c3bd7) --- README.md | 1 + csrc/jit/compiler.hpp | 13 +++++----- .../impls/hybrid_dispatch_unordered.cuh | 25 +++++++++++++++++-- 3 files changed, 31 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index dadda17f4..582c36a1b 100644 --- a/README.md +++ b/README.md @@ -340,6 +340,7 @@ The library provides some environment variables, which may be useful: - `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 diff --git a/csrc/jit/compiler.hpp b/csrc/jit/compiler.hpp index d471e736e..433a1a6ca 100644 --- a/csrc/jit/compiler.hpp +++ b/csrc/jit/compiler.hpp @@ -71,12 +71,13 @@ class Compiler { 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); - // 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"}) + // 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); } diff --git a/deep_ep/include/deep_ep/impls/hybrid_dispatch_unordered.cuh b/deep_ep/include/deep_ep/impls/hybrid_dispatch_unordered.cuh index 816f0571d..d98fa7903 100644 --- a/deep_ep/include/deep_ep/impls/hybrid_dispatch_unordered.cuh +++ b/deep_ep/include/deep_ep/impls/hybrid_dispatch_unordered.cuh @@ -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 __device__ __host__ __forceinline__ int num_sub_parts_at(const int& part_tokens) { if constexpr (kNumSubParts <= 1) { @@ -140,9 +154,16 @@ template 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,