feat(jit): forward sub-part geometry env vars to the JIT

[whn09]

Re-post of Xuan-1998/DeepEP#42 against this repo, rebased onto main (ec623f3) and re-measured on that tree — the old PR's numbers were taken on the pre-refactor tree, so they are not quoted here.

What

hybrid_dispatch_unordered.cuh gates its sub-part geometry behind #ifndef:

macro	default
EP_NUM_SUB_PARTS	2
EP_MIN_SUB_TOKENS	1
EP_SM100_MIN_SUB_TOKENS	15

Nothing in the tree ever sets those macros, so the only way to try a different split today is to edit the header and reinstall. This forwards the three names as JIT -D flags, following the EP_NUM_TOPK_IDX_BITS block immediately above it (and that block's EP_JIT_EXTRA_FLAGS TODO).

Unset ⇒ byte-identical behaviour.

Why it is safe

• All three are device-only. No host translation unit reads them, so a JIT-only define cannot desync host-side sizing from device-side sizing. (This is the reason the patch stops at these three and does not forward, say, kNumMaxTokensPerRank.)
• flags is part of kernel_signature, so changing the env re-JITs instead of silently serving a cached cubin compiled with a different geometry.

Verified end-to-end

EP_JIT_PRINT_COMPILER_COMMAND=1 EP_NUM_SUB_PARTS=1 puts -DEP_NUM_SUB_PARTS=1 on the nvcc line, and the resulting cubin lands in a distinct JIT cache entry.

Measured effect of the knob it exposes

tests/elastic/test_ep.py, 2 × p5en.48xlarge (8×H200 + 16 EFA each), EP8×2 = 16 ranks, --hidden=7168 --num-topk=8 --num-experts=256 --num-sms=12 --allow-hybrid-mode=1 --prefer-overlap-with-compute=0 --test-first-only. 3 reps, variants interleaved within each rep, each variant on its own EP_JIT_CACHE_DIR. Mean over all 16 ranks then over reps; ± is stdev across reps. All 48 rounds exited 0.

EP_NUM_SUB_PARTS=1 on its own:

op	baseline (unset)	EP_NUM_SUB_PARTS=1
prefill (8192 tok) dispatch	1665.1 ± 12.4 µs	1631.9 ± 2.8 µs	−2.0%
prefill cached dispatch	1662.8 ± 8.7 µs	1635.7 ± 1.2 µs	−1.6%
prefill combine	3560.8 ± 9.2 µs	3545.8 ± 5.0 µs	−0.4%
decode (128 tok) dispatch	367.0 ± 12.1 µs	373.7 ± 3.6 µs	+1.8%
decode combine	178.1 ± 5.0 µs	178.5 ± 0.3 µs	+0.2%

So on its own the knob is roughly a wash — a small prefill win, a small decode regression. Its value is that it composes: stacked with #2 it takes decode dispatch from 367.0 µs to 166.1 ± 0.4 µs (−54.7%), where that PR alone reaches 239.8 µs (−34.7%). That is the case this PR is really enabling — being able to find such a combination without a rebuild.

This PR is pure plumbing and changes no default, so it is worth taking independently of whether #2 is accepted.

---

Second architecture: B300 / sm_103 over EFA

Setup. 2 × p6-b300 (8×B300 SXM6 + 16 EFA gen-3 each, EFA installer 1.50.0,
efa.ko 3.3.0g, driver 595.91.07), EP8×2 = 16 ranks, NCCL 2.31.2 GIN with
NCCL_GIN_TYPE=5 (EFA-GDA), same shape flags as above. 3 reps per cell.

On this architecture EP_NUM_SUB_PARTS=1 is neutral, and it does not compose
the way it does on H200. Measured on top of #2 (the two are in one image, so the
b300 run isolates the knob stacked, not standalone):

cell	#2 default	#2 + EP_NUM_SUB_PARTS=1
decode dispatch, kernel-only, bf16	176.84 µs	176.27 µs	−0.3%
decode dispatch, kernel-only, fp8	169.18 µs	169.52 µs	+0.2%
decode end-to-end, bf16	392.9 µs	397.0 µs	+1.0%
decode end-to-end, fp8	554.2 µs	570.3 µs	+2.9%
prefill end-to-end, bf16	2432.6 µs	2436.6 µs	+0.2%
prefill end-to-end, fp8	2207.5 µs	2201.6 µs	−0.3%

So the H200 result — that EP_NUM_SUB_PARTS=1 takes stacked decode dispatch from
239.8 µs to 166.1 µs — does not carry over to sm_103, where #2's clamp is
already the whole effect. That is a point in favour of this PR rather than against
it: the right sub-part split is evidently architecture- and shape-dependent, and
today finding it requires editing a header and reinstalling.

A second, stronger safety argument found while measuring on B300

The claim above that flags is part of kernel_signature is not just a nicety —
it is the only thing that makes a geometry change safe to cache. The other half
of that key, code, is the generated wrapper; the impl headers are only
#included, so their content is not hashed, and signature is just the compiler
version. We hit this directly: two images differing only by a header edit, sharing
one EP_JIT_CACHE_DIR, silently reuse each other's cubins and the edit measures as
a no-op.

Changing this geometry through the env, as this PR allows, is therefore strictly
safer than the header-edit-and-reinstall path it replaces: the env lands in
flags, which is hashed, so a variant can never be served another variant's
cubin. (Worth a separate issue for the header-hash gap, which is orthogonal to
this PR.)

[whn09]

Rebased onto main @ 02efc268; force-pushed. The change itself is unmodified — git patch-id --stable on the rebased commit matches the original.

The old branch showed a 3,000-line diff over 20 files, which was an artefact, not this change. It was cut from main @ ec623f3 (committed 08-21 00:05 UTC); main was then rewritten and force-pushed (cc55cce, committed 08-21 23:51 UTC, same author date). The four base commits exist in both histories under different SHAs, so ec623f3 became unreachable from main and the merge-base fell back six commits to 01dc3aa, attributing the whole unordered-kernel feature to this PR. Both files this PR touches are byte-identical between ec623f3 and main, so the replay was conflict-free.

Measurement caveat: the numbers in the description were taken on ec623f3, not on the rebased tree. ec623f3 → main is +124/−20 across six files, including csrc/kernels/backend/nccl.cu and csrc/kernels/elastic/combine.hpp, so I would not claim they carry over unchanged. The p5en pair used for them is no longer available to me; happy to re-run if you would like the numbers refreshed on main before merging.

[Xuan-1998]

Can you also update the readme's env var list: https://github.com/amazon-contributing/DeepEP#environment-variables

[whn09]

Done in the amended head — added EP_NUM_SUB_PARTS, EP_MIN_SUB_TOKENS, EP_SM100_MIN_SUB_TOKENS to the General list next to EP_NUM_TOPK_IDX_BITS, each noting what 0/unset means and the header default it falls back to. Force-pushed to ca76bfc.