From ca76bfce2331d36345c1fb0c00a0f2e273bb6b0a Mon Sep 17 00:00:00 2001
From: whn09 <wanghenan09@gmail.com>
Date: Fri, 21 Aug 2026 05:36:10 +0000
Subject: [PATCH] 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;