From c1de7637302163efe20ed56a2fd6fb26ce8d8769 Mon Sep 17 00:00:00 2001
From: morluto <76467478+morluto@users.noreply.github.com>
Date: Tue, 28 Jul 2026 10:47:35 +0000
Subject: [PATCH] Validate communication SM count

---
 moonep/api.py                | 27 ++++++++---
 tests/test_api_validation.py | 86 ++++++++++++++++++++++++++++++++++++
 2 files changed, 107 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_api_validation.py

diff --git a/moonep/api.py b/moonep/api.py
index 8d0cbc7..16ed297 100644
--- a/moonep/api.py
+++ b/moonep/api.py
@@ -241,10 +241,25 @@ def _create_context(
                                   non-negative entries encode src_rank * NvS + offv,
                                   -1 marks empty slots
     """
+    if num_sms is not None:
+        if not isinstance(num_sms, int) or isinstance(num_sms, bool):
+            raise TypeError(
+                f"num_sms must be a positive int, got "
+                f"{type(num_sms).__name__}"
+            )
+        if num_sms <= 0:
+            raise ValueError(f"num_sms must be a positive int, got {num_sms}")
+
+    device = torch.cuda.current_device()
+    max_sms = torch.cuda.get_device_properties(device).multi_processor_count
     if num_sms is None:
-        num_sms = 32
-    assert isinstance(num_sms, int) and num_sms > 0, \
-        f"num_sms must be a positive int, got {num_sms}"
+        num_sms = min(32, max_sms)
+    if num_sms > max_sms:
+        raise ValueError(
+            f"num_sms ({num_sms}) must not exceed the device SM count "
+            f"({max_sms})"
+        )
+
     rank = dist.get_rank(group=group)
     R = num_ep_ranks
     assert R == dist.get_world_size(group=group), (
@@ -252,9 +267,7 @@ def _create_context(
         f"({dist.get_world_size(group=group)})"
     )
     N = S * K
-    device = torch.cuda.current_device()
     dev = f"cuda:{device}"
-    max_sms = torch.cuda.get_device_properties(device).multi_processor_count
     num_sms_dedup = _num_sms_dedup_from_env(max_sms)
 
     epn = E // R
@@ -460,7 +473,9 @@ def __init__(
             E: total routed experts in the EP group; must be divisible by
                 ``num_ep_ranks``.
             num_ep_ranks: EP group size (R); must equal the group world size.
-            num_sms: SM count for the comm kernels; None defaults to 32.
+            num_sms: SM count for the comm kernels; must not exceed the
+                current CUDA device's SM count. None defaults to the smaller
+                of 32 and the device SM count.
             token_padding: each non-empty VM group is padded up to a multiple
                 of this token count.
             B: weight prefetch slots per rank; None defaults to
diff --git a/tests/test_api_validation.py b/tests/test_api_validation.py
new file mode 100644
index 0000000..de7efd4
--- /dev/null
+++ b/tests/test_api_validation.py
@@ -0,0 +1,86 @@
+from types import SimpleNamespace
+
+import pytest
+
+from moonep import api
+
+
+@pytest.mark.parametrize(
+    ("num_sms", "error", "match"),
+    [
+        (0, ValueError, "positive int"),
+        (-1, ValueError, "positive int"),
+        (1.5, TypeError, "float"),
+        (True, TypeError, "bool"),
+    ],
+)
+def test_create_context_rejects_invalid_num_sms(
+    monkeypatch, num_sms, error, match
+):
+    monkeypatch.setattr(
+        api.torch.cuda,
+        "current_device",
+        lambda: pytest.fail("invalid num_sms should fail before CUDA access"),
+    )
+
+    with pytest.raises(error, match=match):
+        api._create_context(1, 1, 1, 1, 1, num_sms=num_sms)
+
+
+def test_create_context_rejects_num_sms_above_device_count(monkeypatch):
+    monkeypatch.setattr(api.torch.cuda, "current_device", lambda: 3)
+    monkeypatch.setattr(
+        api.torch.cuda,
+        "get_device_properties",
+        lambda device: (
+            SimpleNamespace(multi_processor_count=8)
+            if device == 3
+            else pytest.fail(f"unexpected device: {device}")
+        ),
+    )
+
+    with pytest.raises(
+        ValueError,
+        match=r"num_sms \(9\) must not exceed the device SM count \(8\)",
+    ):
+        api._create_context(1, 1, 1, 1, 1, num_sms=9)
+
+
+def test_create_context_allows_num_sms_equal_to_device_count(monkeypatch):
+    class ValidationPassed(Exception):
+        pass
+
+    monkeypatch.setattr(api.torch.cuda, "current_device", lambda: 0)
+    monkeypatch.setattr(
+        api.torch.cuda,
+        "get_device_properties",
+        lambda _device: SimpleNamespace(multi_processor_count=8),
+    )
+    monkeypatch.setattr(
+        api.dist,
+        "get_rank",
+        lambda **_kwargs: (_ for _ in ()).throw(ValidationPassed),
+    )
+
+    with pytest.raises(ValidationPassed):
+        api._create_context(1, 1, 1, 1, 1, num_sms=8)
+
+
+def test_create_context_allows_default_on_low_sm_device(monkeypatch):
+    class ValidationPassed(Exception):
+        pass
+
+    monkeypatch.setattr(api.torch.cuda, "current_device", lambda: 0)
+    monkeypatch.setattr(
+        api.torch.cuda,
+        "get_device_properties",
+        lambda _device: SimpleNamespace(multi_processor_count=8),
+    )
+    monkeypatch.setattr(
+        api.dist,
+        "get_rank",
+        lambda **_kwargs: (_ for _ in ()).throw(ValidationPassed),
+    )
+
+    with pytest.raises(ValidationPassed):
+        api._create_context(1, 1, 1, 1, 1)