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

Three bugs in the test and benchmark for
fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert, each of which makes them
report something untrue. Scoped to this operator's own files.

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
   this whole file on hardware that supports the dtype. MetaX C550 reports
   (8, 0) and converts fp8e4nv bit-identically to torch. The check now consults
   an explicit vendor list before falling back to the NVIDIA capability rule.
   Four other files carry the same local copy of this check and have the same
   bug; they are out of scope here and left untouched.

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 its 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. The provider is not always vLLM's own: on
   MetaX it is mcoplib._C, while the vllm wheel there fails to load at all.
   Importing the top-level package is not enough; the compiled submodule must
   be imported before the schemas register.

Author: cheersluvs

benchmark/test_fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 import dataclasses
+import importlib
+import logging
 import random
 
 import pytest
@@ -23,15 +25,67 @@
 
 from . import base
 
+logger = logging.getLogger(__name__)
+
+
+# NVIDIA gates FP8 E4M3 on sm_89+, and this check has historically been spelled
+# `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 skips this whole file
+# on hardware that supports the dtype. Add a vendor here only after verifying
+# on it that a Triton `tl.float8e4nv` conversion matches `torch.float8_e4m3fn`
+# bit-for-bit -- MetaX C550 reports (8, 0) and does (verified 2026-08-05).
+_FP8E4NV_CAPABLE_VENDORS = frozenset({"metax"})
+
 
 def is_support_fp8e4nv():
+    if not hasattr(torch, "float8_e4m3fn"):
+        return False
+    if flag_gems.vendor_name in _FP8E4NV_CAPABLE_VENDORS:
+        return True
     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 registered under `torch.ops._C` by whichever compiled op
+# library the platform ships, and probing that namespace does not import it --
+# so where the reference IS installed, a bare hasattr() reports it missing and
+# no comparison runs. The provider is not always vLLM's own: on MetaX it is
+# `mcoplib._C`, while the `vllm` wheel there fails to load. Note importing the
+# top-level package is not enough; the compiled submodule must be imported
+# before the schemas register.
+VENDOR_OP_LIBS = ("vllm._C", "mcoplib._C")
+
+
+def _load_vendor_ref(op_name):
+    """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` 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 VENDOR_OP_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(VENDOR_OP_LIBS)
+    )
+    return None
+
+
+_VENDOR_REF = _load_vendor_ref("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 +113,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,11 +224,11 @@ 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"
-)
-@pytest.mark.skipif(
-    not is_support_fp8e4nv(), reason="Do not support fp8e4nv when capability < 89"
+    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="Device does not support fp8e4nv")
 def test_fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert():
     bench = FusedDeepseekV4QnormRopeKVRopeQuantInsertBenchmark()
     bench.run()

tests/test_fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.py

@@ -12,16 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import gc
+
 import pytest
 import torch
 
 import flag_gems
+from flag_gems.fused.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert import (
+    fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert as _generic_impl,
+)
 from flag_gems.utils.device_info import get_device_capability
 
 from .conftest import QUICK_MODE
 
+# NVIDIA gates FP8 E4M3 on sm_89+, and this check has historically been spelled
+# `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 skips this whole file
+# on hardware that supports the dtype. Add a vendor here only after verifying
+# on it that a Triton `tl.float8e4nv` conversion matches `torch.float8_e4m3fn`
+# bit-for-bit -- MetaX C550 reports (8, 0) and does (verified 2026-08-05).
+_FP8E4NV_CAPABLE_VENDORS = frozenset({"metax"})
+
 
 def is_support_fp8e4nv():
+    if not hasattr(torch, "float8_e4m3fn"):
+        return False
+    if flag_gems.vendor_name in _FP8E4NV_CAPABLE_VENDORS:
+        return True
     major, minor = get_device_capability()
     return major * 10 + minor >= 89
 
@@ -39,6 +57,71 @@ def is_support_fp8e4nv():
 SCALE_BYTES_PER_TOKEN = NUM_QUANT_BLOCKS + 1  # 8
 HEAD_BYTES = TOKEN_DATA_BYTES + SCALE_BYTES_PER_TOKEN  # 584
 
+# Bytes of device memory the fp32 reference needs per element of
+# [num_tokens, n_heads, HEAD_DIM]. The op itself is cheap; the reference is what
+# blows up, because it upcasts q to fp32: q (bf16, 2B) + q_ref (bf16, 2B) + the
+# fp32 upcast (4B) + the fp32 result (4B) = 12 B/elem live at once, plus
+# allocator headroom. Two observations bracket the real figure:
+#   > 13.2 -- 98304 x 128 OOMs on an 80GB H800 (the shape the old hardcoded
+#             exclusion list was written for)
+#   <=15.5 -- 65536 x 128, an identical element count, passes on a 63.59GB
+#             MetaX C550
+REF_BYTES_PER_ELEM = 14
+
+
+def _free_device_memory():
+    """Free device bytes, or None if the backend cannot report it.
+
+    Two corrections matter, and both were found by getting them wrong on a
+    MetaX C550:
+
+    1. Collect before measuring. The previous parametrized case's tensors are
+       still reachable from its frame, so without this a 64GB card reports
+       ~38GB free and the guard skips shapes that in fact pass.
+    2. Count PyTorch's cached-but-unused blocks as available instead of calling
+       `empty_cache` to release them. Releasing hands them back to the driver,
+       and a marginal allocation then has to be re-served from possibly
+       fragmented driver memory rather than reusing blocks already in hand --
+       which is enough to turn 65536 x 128, a shape that otherwise passes,
+       into an OutOfMemoryError.
+    """
+    try:
+        gc.collect()  # drop the previous case's tensors before measuring
+        driver_free = torch.cuda.mem_get_info()[0]
+        cached = torch.cuda.memory_reserved() - torch.cuda.memory_allocated()
+        return driver_free + cached
+    except Exception:
+        pass
+    try:
+        return torch.cuda.get_device_properties(0).total_memory
+    except Exception:
+        return None
+
+
+def _skip_if_reference_wont_fit(num_tokens: int, n_heads: int):
+    """Skip shapes whose fp32 reference exceeds device memory.
+
+    Replaces a hardcoded exclusion list tuned for an 80GB H800, which still
+    OOMs on 64GB cards -- MetaX C550 died at num_tokens=131072, n_heads=64,
+    inside the reference and before the op was ever called. At
+    REF_BYTES_PER_ELEM the rule reproduces the old H800 exclusion list exactly
+    on an 80GB card, and on a 64GB C550 leaves 131072 x 64 to the
+    OutOfMemoryError handler below rather than pre-skipping it.
+
+    This is a cheap pre-check that avoids allocating tens of GiB only to fail.
+    It is deliberately not the whole story: whether a marginal shape fits also
+    depends on allocator fragmentation (65536 x 128 passes on a C550 while
+    131072 x 64, an identical element count, does not), so the reference call
+    itself also catches OutOfMemoryError.
+    """
+    needed = num_tokens * n_heads * HEAD_DIM * REF_BYTES_PER_ELEM
+    free = _free_device_memory()
+    if free is not None and needed > free:
+        pytest.skip(
+            f"reference needs ~{needed / 2**30:.0f}GiB for num_tokens={num_tokens}, "
+            f"n_heads={n_heads}; only {free / 2**30:.0f}GiB free"
+        )
+
 
 # ─── pytorch reference implementation from vllm ───
 
@@ -266,9 +349,7 @@ def fused_impl(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs):
 
 
 @pytest.mark.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert
-@pytest.mark.skipif(
-    not is_support_fp8e4nv(), reason="Do not support fp8e4nv when capability < 89"
-)
+@pytest.mark.skipif(not is_support_fp8e4nv(), reason="Device does not support fp8e4nv")
 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 64])
 @pytest.mark.parametrize("n_heads", [8, 64])
 def test_q_path_matches_reference(num_tokens: int, n_heads: int):
@@ -326,9 +407,7 @@ def _ue8m0_per_block_scales(kv_roped_nope_f32: torch.Tensor, qblock: int):
 
 
 @pytest.mark.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert
-@pytest.mark.skipif(
-    not is_support_fp8e4nv(), reason="Do not support fp8e4nv when capability < 89"
-)
+@pytest.mark.skipif(not is_support_fp8e4nv(), reason="Device does not support fp8e4nv")
 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 64])
 @pytest.mark.parametrize("block_size", [16, 64])
 def test_kv_path_matches_reference(num_tokens: int, block_size: int):
@@ -374,9 +453,7 @@ def test_kv_path_matches_reference(num_tokens: int, block_size: int):
 
 
 @pytest.mark.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert
-@pytest.mark.skipif(
-    not is_support_fp8e4nv(), reason="Do not support fp8e4nv when capability < 89"
-)
+@pytest.mark.skipif(not is_support_fp8e4nv(), reason="Device does not support fp8e4nv")
 @pytest.mark.parametrize("num_tokens", [4, 17])
 @pytest.mark.parametrize("pad", [1, 5])
 @pytest.mark.parametrize("block_size", [16, 64])
@@ -426,19 +503,15 @@ def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):
 
 
 @pytest.mark.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert
-@pytest.mark.skipif(
-    not is_support_fp8e4nv(), reason="Do not support fp8e4nv when capability < 89"
-)
+@pytest.mark.skipif(not is_support_fp8e4nv(), reason="Device does not support fp8e4nv")
 @pytest.mark.parametrize(
     "num_tokens",
     [1, 4, 17, 64] if QUICK_MODE else [1, 4, 17, 64, 8192, 32768, 65536, 98304, 131072],
 )
 @pytest.mark.parametrize("n_heads", [64, 128])
 @pytest.mark.parametrize("block_size", [16, 64])
 def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
-    # out of memory for huge shape on H800
-    if (num_tokens == 98304 or num_tokens == 131072) and n_heads == 128:
-        return
+    _skip_if_reference_wont_fit(num_tokens, n_heads)
 
     torch.manual_seed(2)
     device = "cuda"
@@ -462,16 +535,37 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
     cos_sin_cache_ref = cos_sin_cache.clone()
     slot_mapping_ref = slot_mapping.clone()
 
-    ref_impl(
-        q_ref,
-        kv_ref,
-        k_cache_ref,
-        slot_mapping_ref,
-        positions_ref,
-        cos_sin_cache_ref,
-        eps,
-        block_size,
-    )
+    try:
+        ref_impl(
+            q_ref,
+            kv_ref,
+            k_cache_ref,
+            slot_mapping_ref,
+            positions_ref,
+            cos_sin_cache_ref,
+            eps,
+            block_size,
+        )
+    except torch.OutOfMemoryError:
+        # The reference is scaffolding, not the code under test, so its running
+        # out of room says nothing about the op -- skip rather than fail.
+        #
+        # Dropping the tensors here is what keeps this from cascading. A failed
+        # case's traceback is retained by pytest, the traceback references this
+        # frame, and the frame keeps q/q_ref/... alive for the rest of the
+        # session: one OOM at num_tokens=131072 pinned 32GiB on a 64GB C550 and
+        # took out an unrelated case 20 tests later. Rebinding to None empties
+        # the frame slots before the exception is raised, so nothing is held.
+        # (`del` would do the same, but leaves the names undefined for the rest
+        # of the scope as far as static analysis is concerned.)
+        q = kv = k_cache = q_ref = kv_ref = k_cache_ref = None
+        positions = positions_ref = cos_sin_cache = cos_sin_cache_ref = None
+        slot_mapping = slot_mapping_ref = None
+        torch.cuda.empty_cache()
+        pytest.skip(
+            f"reference ran out of memory at num_tokens={num_tokens}, "
+            f"n_heads={n_heads}"
+        )
     fused_impl(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, block_size)
 
     torch.testing.assert_close(q, q_ref, rtol=1e-2, atol=1e-2)