[KMCompiler][Benchmark][Test] Fix the fp8e4nv gate, memory guard and vendor op discovery for fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert

Three issues in this operator's test and benchmark, each of which makes the
suite report something untrue.

1. is_support_fp8e4nv() gated on get_device_capability() >= (8, 9). That
   threshold means "Ada or newer" on NVIDIA only; other vendors report their
   own major/minor on a different scale, so it is a false negative that skips
   the whole file on hardware that supports the dtype. MetaX C550 reports
   (8, 0) and converts fp8e4nv bit-identically to torch. Added
   supports_fp8e4nv() to flag_gems/utils/device_info.py, which keeps the
   NVIDIA rule where it means something and lets other vendors opt in
   explicitly. Four other files carry the same local copy of this check and
   are left alone here.

2. The memory guard excluded shapes via a list hardcoded for an 80GB H800 --
   which still OOMs on 64GB cards -- and skipped them with a bare return, so
   they were counted as passed. Replaced with a budget measured against free
   device memory, skipping explicitly. It collects before measuring and counts
   the allocator's cached blocks as available: releasing them with
   empty_cache() instead makes a marginal allocation fail that otherwise
   succeeds.

3. The benchmark probed torch.ops._C for a compiled reference without
   importing the library that registers it, and torch.ops._C gives no hint
   that nothing did. Where the reference IS installed the benchmark reported
   it missing and skipped, so no comparison ran at all. The provider is not
   always vLLM's own: on MetaX it is mcoplib._C, while the vllm wheel there
   fails to load. Added benchmark/vendor_ops.py to try each in turn and log
   why one is rejected.

Author: cheersluvs

benchmark/test_fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.py

@@ -19,19 +19,16 @@
 import torch
 
 import flag_gems
-from flag_gems.utils.device_info import get_device_capability
+from flag_gems.utils.device_info import supports_fp8e4nv as is_support_fp8e4nv
 
 from . import base
+from .vendor_ops import VENDOR_OP_LIBS, load_vendor_op
 
 
-def is_support_fp8e4nv():
-    major, minor = get_device_capability()
-    return major * 10 + minor >= 89
-
-
-VLLM_REF_AVAILABLE = hasattr(
-    torch.ops._C, "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert"
-)
+# The reference is whichever compiled library the platform ships, not
+# necessarily vLLM's own; see benchmark/vendor_ops.py.
+_VENDOR_REF = load_vendor_op("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert")
+VLLM_REF_AVAILABLE = _VENDOR_REF is not None
 HEAD_DIM = 512
 ROPE_DIM = 64
 HEAD_BYTES = 584
@@ -59,7 +56,7 @@ class FusedDeepseekV4QnormRopeKVRopeQuantInsertBenchmark(base.Benchmark):
     def __init__(self):
         super().__init__(
             "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert",
-            torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert,
+            _VENDOR_REF,
             [torch.bfloat16],
         )
         self.set_gems(flag_gems.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)
@@ -170,10 +167,12 @@ def make_input(param: TestParam):
 
 @pytest.mark.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert
 @pytest.mark.skipif(
-    not VLLM_REF_AVAILABLE, reason="The referenced vLLM implementation is not installed"
+    not VLLM_REF_AVAILABLE,
+    reason="No vendor kernel found for this operator (tried %s)"
+    % ", ".join(VENDOR_OP_LIBS),
 )
 @pytest.mark.skipif(
-    not is_support_fp8e4nv(), reason="Do not support fp8e4nv when capability < 89"
+    not is_support_fp8e4nv(), reason="Device does not support fp8e4nv"
 )
 def test_fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert():
     bench = FusedDeepseekV4QnormRopeKVRopeQuantInsertBenchmark()

benchmark/vendor_ops.py

@@ -0,0 +1,69 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Locate vendor-provided reference kernels under `torch.ops._C`.
+
+Several benchmarks compare against a compiled reference registered in the `_C`
+namespace and probe for it with `hasattr(torch.ops._C, name)`. That probe only
+works if something already imported the library that registers the op, and
+`torch.ops._C` gives no hint that nothing did -- so on a platform where the
+reference IS installed the benchmark silently reports it as missing and either
+skips or falls back to an eager-PyTorch baseline, producing speedup numbers
+against the wrong thing.
+
+The registering library is not always vLLM's own. On MetaX the provider is
+`mcoplib._C`, while the `vllm` wheel there is a stock CUDA build whose `_C`
+fails to load at all. Importing only `vllm._C` therefore finds nothing.
+
+Note `import vllm` / `import mcoplib` alone is not enough: the compiled
+submodule has to be imported before the schemas register.
+"""
+
+import importlib
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+# Tried in order; the first one providing the operator wins.
+VENDOR_OP_LIBS = ("vllm._C", "mcoplib._C")
+
+
+def load_vendor_op(op_name, libs=VENDOR_OP_LIBS):
+    """Return `torch.ops._C.<op_name>`, importing vendor libraries as needed.
+
+    Returns None if no library provides it. Import failures are logged rather
+    than swallowed -- a silent `except: pass` here is what makes a missing
+    baseline indistinguishable from an unimportable one.
+    """
+    fn = getattr(torch.ops._C, op_name, None)
+    if callable(fn):
+        return fn
+
+    for lib in libs:
+        try:
+            importlib.import_module(lib)
+        except Exception as e:
+            logger.info("vendor op library %s unavailable: %s", lib, e)
+            continue
+        fn = getattr(torch.ops._C, op_name, None)
+        if callable(fn):
+            logger.info("found %s in %s", op_name, lib)
+            return fn
+        logger.info("%s loaded but does not provide %s", lib, op_name)
+
+    logger.info(
+        "no vendor kernel for %s (tried %s)", op_name, ", ".join(libs)
+    )
+    return None

src/flag_gems/utils/device_info.py

@@ -16,6 +16,9 @@
 from dataclasses import dataclass
 from functools import lru_cache
 
+import torch
+
+from flag_gems.runtime import device as _runtime_device
 from flag_gems.runtime import torch_device_fn
 
 
@@ -69,6 +72,36 @@ def get_device_capability() -> tuple[int, int]:
         return (0, 0)
 
 
+# Vendors whose hardware does FP8 E4M3 but whose reported compute capability is
+# below the NVIDIA sm_89 threshold, so the capability rule would wrongly exclude
+# them. Add a vendor here only after verifying on the hardware that a Triton
+# `tl.float8e4nv` conversion matches `torch.float8_e4m3fn` bit-for-bit.
+#   metax: C550 reports capability (8, 0) and converts bit-for-bit identically
+#   to torch (verified 2026-08-05 on MetaX C550).
+_FP8E4NV_CAPABLE_VENDORS = frozenset({"metax"})
+
+
+@lru_cache(maxsize=1)
+def supports_fp8e4nv() -> bool:
+    """Whether this device supports FP8 E4M3 (`tl.float8e4nv`).
+
+    NVIDIA gates E4M3 on sm_89+, and tests have historically spelled that check
+    as `get_device_capability() >= (8, 9)`. That number only means "Ada or
+    newer" on NVIDIA -- other vendors report their own major/minor on a
+    different scale, so applying the threshold to them is a false negative that
+    silently skips whole test files on hardware that supports the dtype.
+
+    Keep the capability rule where it is meaningful, and consult an explicit
+    per-vendor list elsewhere.
+    """
+    if not hasattr(torch, "float8_e4m3fn"):
+        return False
+    if _runtime_device.vendor_name in _FP8E4NV_CAPABLE_VENDORS:
+        return True
+    major, minor = get_device_capability()
+    return major * 10 + minor >= 89
+
+
 @lru_cache(maxsize=1)
 def get_device_info() -> DeviceInfo:
     props = get_device_properties()