diff --git a/vllm_fl/dispatch/backends/flaggems/flaggems.py b/vllm_fl/dispatch/backends/flaggems/flaggems.py index df5e3d58..04808b09 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 ff0d684d..36e1eaaa 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 02128a6d..81f35402 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(