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.
Everything about which accelerator a run targets was previously spelled CUDA:
device discovery went through nvidia-smi, pinning through CUDA_VISIBLE_DEVICES,
the vendor guard only recognised NVIDIA, and the -p plugin list was hardcoded in
run_pytest.sh with no way to drop the CUDA-only members. Running the same
operator corpus on another vendor's device therefore meant editing the drivers.
The selection now hangs off one ZL_BACKEND variable (cuda/musa/rocm/ascend, with
the usual aliases):
- backends/ owns per-backend device enumeration, visibility env keys and
capability flags, so nothing above it names a vendor;
- plugin_plan.py derives the -p list from those capabilities instead of a
hardcoded array, which is what lets a non-CUDA run omit the graph-capture and
cluster-fix plugins rather than watch them fail;
- _timing_plugin selects the timer from the same capabilities: CUDA keeps
graph capture with its per-shape fallback, other backends stay on upstream
do_bench and record timer=direct, because capture-then-fallback on a device
without graph support only produces noise;
- _device_setup_plugin replaces the NVIDIA-only vendor guard;
- providers/flaggems.py resolves the pytest node, so a caller may pass TEST_NODE
directly instead of the OP/OP_FILE pair, and SHAPE_MODE=upstream can defer to
FlagGems' own core_shapes;
- _run_identity.py writes run_identity.json next to each run, recording backend,
device, timer, autotune mode and the resolved test node. Two runs whose
calibers differ are otherwise indistinguishable after the fact.
run_batch gains --backend and --devices (--gpus kept as an alias) and carries
backend/device/timer into ops_status.csv. test_backends.py covers the registry,
alias resolution, visibility env and plugin planning.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Aharrypotter
