From ff1c8e427f165edeaecc90747e9dd954ac75be11 Mon Sep 17 00:00:00 2001
From: Qiming Teng <tengqm@outlook.com>
Date: Fri, 31 Jul 2026 16:07:52 +0800
Subject: [PATCH] feat(metax): route reshape_and_cache_flash through flag_gems
 dispatch

On an empty vLLM wheel (VLLM_TARGET_DEVICE=empty) there is no compiled
_C_cache_ops kernel, so the MetaX attention path crashes at the first
forward pass: fa_utils.py binds reshape_and_cache_flash to
vllm._custom_ops, which delegates to the missing
torch.ops._C_cache_ops.reshape_and_cache_flash and raises AttributeError.

Register reshape_and_cache_flash as a flag_gems op (default.flagos,
guarded by _has_flaggems_op) backed by the pure-Triton
flag_gems.fused.reshape_and_cache_flash, and bind the metax call site to
CachedOp("reshape_and_cache_flash"). No coupling to vllm's private
_C_cache_ops ABI; standard wheels keep their existing selection.

This PR was written in part with the assistance of generative AI.
---
 vllm_fl/dispatch/backends/flaggems/flaggems.py     | 14 ++++++++++++++
 vllm_fl/dispatch/backends/flaggems/register_ops.py | 14 ++++++++++++++
 .../vendor/metax/impl/attention/utils/fa_utils.py  | 10 +++++++++-
 3 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/vllm_fl/dispatch/backends/flaggems/flaggems.py b/vllm_fl/dispatch/backends/flaggems/flaggems.py
index df5e3d587..04808b090 100644
--- a/vllm_fl/dispatch/backends/flaggems/flaggems.py
+++ b/vllm_fl/dispatch/backends/flaggems/flaggems.py
@@ -121,6 +121,20 @@ def silu_and_mul_with_clamp(self, x: torch.Tensor, swiglu_limit: float, swiglu_l
 
         return silu_and_mul_with_clamp_flaggems(x, swiglu_limit_tensor)
 
+    def reshape_and_cache_flash(self, *args, **kwargs):
+        """Write key/value into the paged KV cache (flash layout).
+
+        Forwards to flag_gems.fused.reshape_and_cache_flash, which is a
+        pure-Triton in-place op. Sourced from the concrete submodule (not the
+        lazy top-level flag_gems re-export) so it resolves to the real
+        implementation regardless of flag_gems init ordering. Call signature:
+        (key, value, key_cache, value_cache, slot_mapping, kv_cache_dtype,
+        k_scale, v_scale) -> None.
+        """
+        from flag_gems.fused import reshape_and_cache_flash
+
+        return reshape_and_cache_flash(*args, **kwargs)
+
     def rms_norm(
         self,
         obj,
diff --git a/vllm_fl/dispatch/backends/flaggems/register_ops.py b/vllm_fl/dispatch/backends/flaggems/register_ops.py
index ff0d684d3..36e1eaaa7 100644
--- a/vllm_fl/dispatch/backends/flaggems/register_ops.py
+++ b/vllm_fl/dispatch/backends/flaggems/register_ops.py
@@ -121,6 +121,20 @@ def register_builtins(registry) -> None:
             vendor=None,
             priority=BackendPriority.DEFAULT,
         ),
+        # reshape_and_cache_flash (KV cache write, flash layout).
+        # Needed on empty vllm wheels (VLLM_TARGET_DEVICE=empty) that ship no
+        # compiled _C_cache_ops kernel; flag_gems provides a pure-Triton impl.
+        OpImpl(
+            op_name="reshape_and_cache_flash",
+            impl_id="default.flagos",
+            kind=BackendImplKind.DEFAULT,
+            fn=_bind_is_available(
+                backend.reshape_and_cache_flash,
+                _has_flaggems_op("reshape_and_cache_flash"),
+            ),
+            vendor=None,
+            priority=BackendPriority.DEFAULT,
+        ),
         # fused topk bias
         OpImpl(
             op_name="fused_topk_bias",
diff --git a/vllm_fl/dispatch/backends/vendor/metax/impl/attention/utils/fa_utils.py b/vllm_fl/dispatch/backends/vendor/metax/impl/attention/utils/fa_utils.py
index 02128a6dd..81f35402b 100644
--- a/vllm_fl/dispatch/backends/vendor/metax/impl/attention/utils/fa_utils.py
+++ b/vllm_fl/dispatch/backends/vendor/metax/impl/attention/utils/fa_utils.py
@@ -22,7 +22,15 @@ def get_scheduler_metadata(*args, **kwargs) -> None:
             return None
 
     get_scheduler_metadata = _dummy_ops.get_scheduler_metadata
-    reshape_and_cache_flash = ops.reshape_and_cache_flash
+    # Route reshape_and_cache_flash through the FL dispatch manager instead of
+    # binding vllm._custom_ops.reshape_and_cache_flash directly. On an empty
+    # vllm wheel (VLLM_TARGET_DEVICE=empty) the latter dispatches to the missing
+    # torch.ops._C_cache_ops C kernel and raises at the first forward pass.
+    # CachedOp resolves to the flag_gems (Triton) implementation registered
+    # under op_name "reshape_and_cache_flash", with policy-driven fallback.
+    from vllm_fl.dispatch import CachedOp
+
+    reshape_and_cache_flash = CachedOp("reshape_and_cache_flash")
 
 
 def get_flash_attn_version(