feat(metax): route reshape_and_cache_flash through flag_gems dispatch

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,

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",

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(