bench: report when a foreign vendor's FlagGems implementation gets loaded

FlagGems keeps per-vendor implementations under
flag_gems/runtime/backend/_<vendor>/ops/, and benchmarks commonly gate on
whether one imports:

    try:
        from flag_gems.runtime.backend._nvidia.hopper.ops.mm import router_gemm
        ROUTER_GEMM_AVAILABLE = True
    except Exception:
        ROUTER_GEMM_AVAILABLE = False

An import only proves the .py is on disk. It touches no device, so it succeeds
on every backend and the gate never fires. test_router_gemm.py therefore ran
the NVIDIA Hopper kernel on an S5000: its autotune and heuristics are keyed to
nvidia, supply no BLOCK_M/N/K/SPLIT_K on musa, and the run died deep inside the
JIT with "dynamic_func() missing 4 required positional arguments" — an error
that points nowhere near the cause. Meanwhile _mthreads/ops/mm.py holds a
working router_gemm the benchmark never imports. Gating on vendor_name first,
the way the Cambricon branch in the same file already does, makes it run: 10
shapes, 10 kernels, cvt 10 -> 2.

The plugin now scans sys.modules at session end and prints
FOREIGN_VENDOR_MODULE with the vendors it found. It reports rather than
refuses: loading another vendor's module is not proof of a problem, but when a
result looks wrong it is the first line to check. Vendor roots are compared
segment-wise, not by prefix, so a vendor whose name extends ours is not
mistaken for our own.

README gains the ordered checklist for an operator that will not run, with the
real instance behind each entry.
This commit is contained in:
Aharrypotter committed 2026-09-10 18:51:26 +08:00
1 parent 2e456bdc74
commit 4ac2a9acd8
3 files changed
+140

No files matched your search

+45
View File
@@ -614,5 +614,50 @@ class TelemetryTests(unittest.TestCase):
os.environ["MUSA_VISIBLE_DEVICES"] = old
class ForeignVendorModuleTests(unittest.TestCase):
"""benchmark 的可用性判据常写成 `try: import <厂商实现>; AVAILABLE = True`。
import 只证明那个 .py 在磁盘上——它不碰设备,所以在任何后端上都成功,gate 于是
形同虚设。router_gemm 就这样让 NVIDIA Hopper 的 kernel 跑在了 MUSA 卡上,最后
死在 JIT 深处的 "dynamic_func() missing BLOCK_M ...",报错完全指不向真正的原因。
"""
P = "flag_gems.runtime.backend._"
def _mods(self, *names):
return {self.P + n: object() for n in names} | {"torch": object()}
def test_only_own_vendor_is_clean(self):
from _device_setup_plugin import foreign_vendor_modules
mods = self._mods("mthreads", "mthreads.ops.mm")
self.assertEqual(foreign_vendor_modules("mthreads", mods), [])
def test_foreign_vendor_is_reported_once_per_root(self):
from _device_setup_plugin import foreign_vendor_modules
# 同一厂商的多个子模块只报一条
mods = self._mods("mthreads.ops.mm",
"nvidia.hopper.ops.mm", "nvidia.ops.mm", "nvidia")
self.assertEqual(foreign_vendor_modules("mthreads", mods), ["nvidia"])
def test_several_foreign_vendors(self):
from _device_setup_plugin import foreign_vendor_modules
mods = self._mods("mthreads", "nvidia.ops.mm", "cambricon.ops.mm")
self.assertEqual(foreign_vendor_modules("mthreads", mods),
["cambricon", "nvidia"])
def test_a_vendor_whose_name_extends_ours_is_still_foreign(self):
# 按前缀比会把 _mthreads_extra 认成自家
from _device_setup_plugin import foreign_vendor_modules
mods = self._mods("mthreads", "mthreads_extra.ops")
self.assertEqual(foreign_vendor_modules("mthreads", mods),
["mthreads_extra"])
def test_unrelated_modules_are_ignored(self):
from _device_setup_plugin import foreign_vendor_modules
mods = {"torch": object(), "triton": object(),
"flag_gems.ops.mm": object()}
self.assertEqual(foreign_vendor_modules("mthreads", mods), [])
if __name__ == "__main__":
unittest.main()