From e0f110a3771293fe493bc44a84e626b1192c1383 Mon Sep 17 00:00:00 2001 From: Xuan Jiang Date: Wed, 29 Jul 2026 04:48:04 +0000 Subject: [PATCH 1/3] envs: probe RDMA link rate via sysfs, survive probe failure get_rdma_gbs() only knew how to ask ibstat for a CA named EP_NIC_NAME (default mlx5_0). EFA devices expose no umad CA, so on EFA hosts the probe returned 0 and get_theoretical_num_sms() divided by it -- any multi-node run without an explicit --num-sms crashed with ZeroDivisionError. Read /sys/class/infiniband//ports/*/rate first, which works for every verbs provider, and fall back to ibstat for setups whose rate only shows there. When EP_NIC_NAME is unset and the default device is absent, pick the fastest device under /sys/class/infiniband instead of failing; an explicitly named NIC still fails loudly rather than guessing. --- deep_ep/utils/envs.py | 55 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/deep_ep/utils/envs.py b/deep_ep/utils/envs.py index f6e34d988..af4183fcb 100644 --- a/deep_ep/utils/envs.py +++ b/deep_ep/utils/envs.py @@ -1,3 +1,14 @@ +# MIT License +# +# Copyright (c) 2025 DeepSeek +# Changes and additions copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + import functools import inspect import os @@ -242,17 +253,61 @@ def check_fast_rdma_atomic_support(nic_name: str = _DEFAULT_NIC_NAME) -> bool: return False +def _get_sysfs_rdma_gbs(nic_name: str) -> float: + """ + Read one RDMA device's link rate from sysfs (`/sys/class/infiniband//ports/*/rate`). + Works for any verbs provider (mlx5, EFA's `rdmap*`, ...) without external tools. + + Arguments: + nic_name: the NIC device name. + + Returns: + gbs: the device's link rate in GB/s (0 if the device or its rate is unavailable). + """ + rate = 0 + ports_dir = os.path.join('/sys/class/infiniband', nic_name, 'ports') + try: + for port in os.listdir(ports_dir): + with open(os.path.join(ports_dir, port, 'rate')) as f: + match = re.match(r'\s*(\d+)\s*Gb/sec', f.read()) + if match: + rate = max(rate, int(match.group(1))) + except OSError: + pass + return rate / 8 + + @functools.lru_cache() def get_rdma_gbs(nic_name: str = _DEFAULT_NIC_NAME) -> float: """ Get the RDMA bandwidth in GB/s, cached. + Probes sysfs first, which covers any verbs provider; `ibstat` is kept as a fallback but + cannot see providers without a umad interface (e.g. EFA). When `EP_NIC_NAME` is not set + and the default device does not exist (EFA hosts have no `mlx5_0`), the fastest device + under `/sys/class/infiniband` is used instead. + Arguments: nic_name: the NIC device name. Returns: gbs: the RDMA bandwidth in GB/s (0 if detection fails). """ + gbs = _get_sysfs_rdma_gbs(nic_name) + if gbs > 0: + return gbs + + # The un-overridden default may simply not exist on this fabric; an explicitly named NIC + # must not fall back silently + if 'EP_NIC_NAME' not in os.environ: + try: + devices = sorted(os.listdir('/sys/class/infiniband')) + except OSError: + devices = [] + gbs = max((_get_sysfs_rdma_gbs(device) for device in devices), default=0) + if gbs > 0: + return gbs + # noinspection PyBroadException try: result = subprocess.run(['ibstat'], capture_output=True, text=True, check=True) From 02efc268a37802fc00812ede8f5ad7f535ceea0e Mon Sep 17 00:00:00 2001 From: Aviv Benchorin Date: Fri, 21 Aug 2026 17:54:06 +0000 Subject: [PATCH 2/3] tests: bounded pressure loops and teardown barrier for test_ep Add --pressure-iterations to bound the pressure-test loop (upstream's --do-pressure-test runs int(1e9) seeds, i.e. until killed; 0 keeps that behavior), with argument validation. Add a barrier before dist.destroy_process_group() so a fast rank cannot tear down the TCPStore while slower ranks are still in destroy. Signed-off-by: Xuan Jiang --- tests/elastic/test_ep.py | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/tests/elastic/test_ep.py b/tests/elastic/test_ep.py index 7344af5f2..456e20340 100644 --- a/tests/elastic/test_ep.py +++ b/tests/elastic/test_ep.py @@ -1,3 +1,14 @@ +# MIT License +# +# Copyright (c) 2025 DeepSeek +# Changes and additions copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + import argparse import os import torch @@ -545,7 +556,12 @@ def construct_elastic_buffer(): test_dispatch_combine(buffer, args) # Pressure tests - for seed in range(int(1e9) if args.do_pressure_test else 0): + if args.do_pressure_test: + pressure_iteration_count = args.pressure_iterations if args.pressure_iterations != 0 else int(1e9) + else: + pressure_iteration_count = 0 + + for seed in range(pressure_iteration_count): if not args.reuse_elastic_buffer: # Recreate elastic buffer buffer.destroy() @@ -558,6 +574,7 @@ def construct_elastic_buffer(): # Destroy the runtime and communication group buffer.destroy() + dist.barrier() dist.destroy_process_group() @@ -591,6 +608,12 @@ def construct_elastic_buffer(): parser.add_argument('--skip-check', action='store_true', help='Whether to skip correctness checks') parser.add_argument('--skip-perf-test', action='store_true', help='Whether to skip performance tests') parser.add_argument('--do-pressure-test', action='store_true', help='Whether to do pressure test') + parser.add_argument( + '--pressure-iterations', + type=int, + default=0, + help='Number of pressure-loop seeds; 0 represents the default unbounded value of 1e9 seeds', + ) parser.add_argument('--reuse-elastic-buffer', action='store_true', help='Whether to reuse elastic buffer for each test') parser.add_argument('--test-first-only', action='store_true', help='Only test the first case') parser.add_argument('--unbalanced-ratio', type=float, default=1.0, help='The MoE unbalanced ratio') @@ -599,6 +622,10 @@ def construct_elastic_buffer(): parser.add_argument('--dump-profile-traces', type=str, default='', help='Dump profiling trace JSONs') parser.add_argument('--ignore-local-traffic', action='store_true', help='Whether to ignore local traffic during bandwidth calculation') args = parser.parse_args() + if args.pressure_iterations < 0: + parser.error("--pressure-iterations must be non-negative") + if args.pressure_iterations and not args.do_pressure_test: + parser.error("--pressure-iterations requires --do-pressure-test") # Create dump trace directories if args.dump_profile_traces: From 5118c2e3f97920841567eafd868b850a091ff3a0 Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 21 Aug 2026 05:36:10 +0000 Subject: [PATCH 3/3] 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. --- csrc/jit/compiler.hpp | 9 +++++++++ 1 file changed, 9 insertions(+) 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;