From 15772189282948ff6a2c1561d589e51dfde92c34 Mon Sep 17 00:00:00 2001 From: Mr-Neutr0n Date: Sat, 1 Aug 2026 00:49:56 +0530 Subject: [PATCH 1/3] Handle missing bandwidth in SM estimation --- deep_ep/buffers/ep.py | 15 +++++-- tests/utils/test_elastic_host.py | 72 ++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 3 deletions(-) create mode 100644 tests/utils/test_elastic_host.py diff --git a/deep_ep/buffers/ep.py b/deep_ep/buffers/ep.py index 9b8a95364..9dba7c06c 100644 --- a/deep_ep/buffers/ep.py +++ b/deep_ep/buffers/ep.py @@ -508,8 +508,17 @@ def get_expected_topk(num_groups: int) -> float: nvlink_traffic += self.num_nvlink_ranks / self.num_ranks * (1 - 1 / self.num_nvlink_ranks) # Except local bypass rdma_traffic += (self.num_ranks - self.num_nvlink_ranks) / self.num_ranks - # Found the bounded one - if self.num_scaleout_ranks > 1 and (rdma_traffic / rdma_gbs) > (nvlink_traffic / nvlink_gbs): + # Find the bounded link. Detection helpers return 0 on failure, so treat + # a link with traffic but unknown bandwidth as unbounded transfer time + # instead of dividing by zero. + def get_transfer_time(traffic: float, bandwidth: float) -> float: + if traffic == 0: + return 0 + return traffic / bandwidth if bandwidth > 0 else math.inf + + rdma_time = get_transfer_time(rdma_traffic, rdma_gbs) + nvlink_time = get_transfer_time(nvlink_traffic, nvlink_gbs) + if self.num_scaleout_ranks > 1 and rdma_time > nvlink_time: bounded_traffic, bounded_gbs = rdma_traffic, rdma_gbs else: bounded_traffic, bounded_gbs = nvlink_traffic, nvlink_gbs @@ -518,7 +527,7 @@ def get_expected_topk(num_groups: int) -> float: # NOTES: will try to use more SMs if not overlap with compute num_device_sms = torch.cuda.get_device_properties('cuda').multi_processor_count num_sms = num_device_sms # No traffic, e.g., EP=1 - if bounded_traffic > 0: + if bounded_traffic > 0 and bounded_gbs > 0: num_sms = max( # Dispatch bounded_gbs / bounded_traffic * sm_read / sm_read_gbs, diff --git a/tests/utils/test_elastic_host.py b/tests/utils/test_elastic_host.py new file mode 100644 index 000000000..b47df9f9d --- /dev/null +++ b/tests/utils/test_elastic_host.py @@ -0,0 +1,72 @@ +import importlib.util +import math +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +import torch + + +def add_module(name, **attributes): + module = types.ModuleType(name) + for key, value in attributes.items(): + setattr(module, key, value) + sys.modules[name] = module + return module + + +def load_elastic(): + package = add_module('deep_ep') + package.__path__ = [] + buffers = add_module('deep_ep.buffers') + buffers.__path__ = [] + utils = add_module('deep_ep.utils') + utils.__path__ = [] + + add_module('deep_ep._C', EventHandle=type('EventHandle', (), {})) + add_module('deep_ep.utils.event', EventOverlap=type('EventOverlap', (), {})) + add_module('deep_ep.utils.math', align=lambda x, y: math.ceil(x / y) * y) + add_module( + 'deep_ep.utils.semantic', + value_or=lambda value, default: default if value is None else value, + weak_lru=lambda *args, **kwargs: lambda function: function, + ) + add_module( + 'deep_ep.utils.envs', + check_fast_rdma_atomic_support=lambda: True, + check_nvlink_connections=lambda group: None, + check_torch_deterministic=lambda: None, + get_nvlink_gbs=lambda: 0, + get_rdma_gbs=lambda: 0, + ) + add_module('deep_ep.utils.comm', get_nccl_comm_handle=lambda *args, **kwargs: None) + + path = Path(__file__).parents[2] / 'deep_ep' / 'buffers' / 'elastic.py' + spec = importlib.util.spec_from_file_location('deep_ep.buffers.elastic', path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class TheoreticalSMTests(unittest.TestCase): + def test_missing_hybrid_bandwidth_falls_back_to_device_sm_count(self): + elastic = load_elastic() + buffer = elastic.ElasticBuffer.__new__(elastic.ElasticBuffer) + buffer.num_rdma_ranks = 2 + buffer.num_nvlink_ranks = 8 + buffer.num_scaleout_ranks = 2 + buffer.num_scaleup_ranks = 8 + buffer.num_ranks = 16 + buffer.prefer_overlap_with_compute = True + + properties = types.SimpleNamespace(multi_processor_count=132) + with patch.object(torch.cuda, 'get_device_properties', return_value=properties): + num_sms = buffer.get_theoretical_num_sms(64, 8) + + self.assertEqual(num_sms, 132) + + +if __name__ == '__main__': + unittest.main() From d17262c2485286b5bdc9400cbedf38e80b1b7111 Mon Sep 17 00:00:00 2001 From: Mr-Neutr0n Date: Mon, 3 Aug 2026 06:33:37 +0530 Subject: [PATCH 2/3] test: isolate DeepEP module stubs --- tests/utils/test_elastic_host.py | 71 ++++++++++++++++++++------------ 1 file changed, 45 insertions(+), 26 deletions(-) diff --git a/tests/utils/test_elastic_host.py b/tests/utils/test_elastic_host.py index b47df9f9d..219ce10f5 100644 --- a/tests/utils/test_elastic_host.py +++ b/tests/utils/test_elastic_host.py @@ -18,36 +18,55 @@ def add_module(name, **attributes): def load_elastic(): - package = add_module('deep_ep') - package.__path__ = [] - buffers = add_module('deep_ep.buffers') - buffers.__path__ = [] - utils = add_module('deep_ep.utils') - utils.__path__ = [] - - add_module('deep_ep._C', EventHandle=type('EventHandle', (), {})) - add_module('deep_ep.utils.event', EventOverlap=type('EventOverlap', (), {})) - add_module('deep_ep.utils.math', align=lambda x, y: math.ceil(x / y) * y) - add_module( + module_names = ( + 'deep_ep', + 'deep_ep.buffers', + 'deep_ep._C', + 'deep_ep.utils', + 'deep_ep.utils.event', + 'deep_ep.utils.math', 'deep_ep.utils.semantic', - value_or=lambda value, default: default if value is None else value, - weak_lru=lambda *args, **kwargs: lambda function: function, - ) - add_module( 'deep_ep.utils.envs', - check_fast_rdma_atomic_support=lambda: True, - check_nvlink_connections=lambda group: None, - check_torch_deterministic=lambda: None, - get_nvlink_gbs=lambda: 0, - get_rdma_gbs=lambda: 0, + 'deep_ep.utils.comm', ) - add_module('deep_ep.utils.comm', get_nccl_comm_handle=lambda *args, **kwargs: None) + saved_modules = {name: sys.modules.get(name) for name in module_names} + try: + package = add_module('deep_ep') + package.__path__ = [] + buffers = add_module('deep_ep.buffers') + buffers.__path__ = [] + utils = add_module('deep_ep.utils') + utils.__path__ = [] - path = Path(__file__).parents[2] / 'deep_ep' / 'buffers' / 'elastic.py' - spec = importlib.util.spec_from_file_location('deep_ep.buffers.elastic', path) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module + add_module('deep_ep._C', EventHandle=type('EventHandle', (), {})) + add_module('deep_ep.utils.event', EventOverlap=type('EventOverlap', (), {})) + add_module('deep_ep.utils.math', align=lambda x, y: math.ceil(x / y) * y) + add_module( + 'deep_ep.utils.semantic', + value_or=lambda value, default: default if value is None else value, + weak_lru=lambda *args, **kwargs: lambda function: function, + ) + add_module( + 'deep_ep.utils.envs', + check_fast_rdma_atomic_support=lambda: True, + check_nvlink_connections=lambda group: None, + check_torch_deterministic=lambda: None, + get_nvlink_gbs=lambda: 0, + get_rdma_gbs=lambda: 0, + ) + add_module('deep_ep.utils.comm', get_nccl_comm_handle=lambda *args, **kwargs: None) + + path = Path(__file__).parents[2] / 'deep_ep' / 'buffers' / 'elastic.py' + spec = importlib.util.spec_from_file_location('deep_ep.buffers.elastic', path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + finally: + for name, saved_module in saved_modules.items(): + if saved_module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = saved_module class TheoreticalSMTests(unittest.TestCase): From fd6949e93faa9fc06037ecf7d045b215be03e0c4 Mon Sep 17 00:00:00 2001 From: Mr-Neutr0n Date: Mon, 3 Aug 2026 06:39:14 +0530 Subject: [PATCH 3/3] Address review suggestions for bandwidth fallback --- deep_ep/buffers/ep.py | 4 +++- tests/utils/test_elastic_host.py | 24 +++++++++++++++--------- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/deep_ep/buffers/ep.py b/deep_ep/buffers/ep.py index 9dba7c06c..6c71cdc75 100644 --- a/deep_ep/buffers/ep.py +++ b/deep_ep/buffers/ep.py @@ -513,11 +513,13 @@ def get_expected_topk(num_groups: int) -> float: # instead of dividing by zero. def get_transfer_time(traffic: float, bandwidth: float) -> float: if traffic == 0: - return 0 + return 0.0 return traffic / bandwidth if bandwidth > 0 else math.inf rdma_time = get_transfer_time(rdma_traffic, rdma_gbs) nvlink_time = get_transfer_time(nvlink_traffic, nvlink_gbs) + # If both are unknown, either branch selects zero bandwidth and the + # calculation below conservatively retains the full device SM count. if self.num_scaleout_ranks > 1 and rdma_time > nvlink_time: bounded_traffic, bounded_gbs = rdma_traffic, rdma_gbs else: diff --git a/tests/utils/test_elastic_host.py b/tests/utils/test_elastic_host.py index 219ce10f5..b3b1a63b6 100644 --- a/tests/utils/test_elastic_host.py +++ b/tests/utils/test_elastic_host.py @@ -56,8 +56,8 @@ def load_elastic(): ) add_module('deep_ep.utils.comm', get_nccl_comm_handle=lambda *args, **kwargs: None) - path = Path(__file__).parents[2] / 'deep_ep' / 'buffers' / 'elastic.py' - spec = importlib.util.spec_from_file_location('deep_ep.buffers.elastic', path) + path = Path(__file__).parents[2] / 'deep_ep' / 'buffers' / 'ep.py' + spec = importlib.util.spec_from_file_location('deep_ep.buffers.ep', path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module @@ -69,16 +69,22 @@ def load_elastic(): sys.modules[name] = saved_module +def make_hybrid_buffer(elastic): + """Build the minimal EPBuffer topology used by the host-only test.""" + buffer = elastic.EPBuffer.__new__(elastic.EPBuffer) + buffer.num_rdma_ranks = 2 + buffer.num_nvlink_ranks = 8 + buffer.num_scaleout_ranks = 2 + buffer.num_scaleup_ranks = 8 + buffer.num_ranks = 16 + buffer.prefer_overlap_with_compute = True + return buffer + + class TheoreticalSMTests(unittest.TestCase): def test_missing_hybrid_bandwidth_falls_back_to_device_sm_count(self): elastic = load_elastic() - buffer = elastic.ElasticBuffer.__new__(elastic.ElasticBuffer) - buffer.num_rdma_ranks = 2 - buffer.num_nvlink_ranks = 8 - buffer.num_scaleout_ranks = 2 - buffer.num_scaleup_ranks = 8 - buffer.num_ranks = 16 - buffer.prefer_overlap_with_compute = True + buffer = make_hybrid_buffer(elastic) properties = types.SimpleNamespace(multi_processor_count=132) with patch.object(torch.cuda, 'get_device_properties', return_value=properties):