Handle missing bandwidth in SM estimation

deep_ep/buffers/ep.py

@@ -508,8 +508,19 @@ 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.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:
             bounded_traffic, bounded_gbs = nvlink_traffic, nvlink_gbs
@@ -518,7 +529,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,

tests/utils/test_elastic_host.py

@@ -0,0 +1,97 @@
+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():
+    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',
+        'deep_ep.utils.envs',
+        'deep_ep.utils.comm',
+    )
+    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__ = []
+
+        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' / '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
+    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
+
+
+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 = make_hybrid_buffer(elastic)
+
+        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()