[compile] Fix fake kernel return dtype (#50444)

Signed-off-by: Richard Zou <zou3519@gmail.com>
Co-authored-by: Codex <noreply@openai.com>

tests/compile/passes/test_rope_kvcache_fusion.py

@@ -29,6 +29,7 @@ from vllm.config import (
 from vllm.config.utils import Range
 from vllm.forward_context import get_forward_context, set_forward_context
 from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.attention import attention as attention_module
 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import _encode_layer_name
@@ -43,6 +44,37 @@ VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
 FP8_DTYPE = current_platform.fp8_dtype()
 
 
+def test_kv_cache_update_ops_fake_tensor_metadata(monkeypatch: pytest.MonkeyPatch):
+    query = torch.empty(1, dtype=torch.bfloat16)
+    kv_cache = torch.empty(1, dtype=torch.uint8)
+    context = (None, None, kv_cache, None)
+    monkeypatch.setattr(attention_module, "get_attention_context", lambda _: context)
+    monkeypatch.setattr(rope_kvcache_fusion, "get_attention_context", lambda _: context)
+
+    runtime_output = attention_module.unified_kv_cache_update(query, query, "layer")
+    fake_output = attention_module.unified_kv_cache_update_fake(query, query, "layer")
+    assert runtime_output.shape == fake_output.shape
+    assert runtime_output.dtype == fake_output.dtype
+    assert runtime_output.device == fake_output.device
+
+    args = (
+        query,
+        query,
+        query,
+        torch.empty(1, dtype=torch.int64),
+        torch.empty(1),
+        True,
+        "layer",
+    )
+    runtime_output = rope_kvcache_fusion.fused_rope_and_unified_kv_cache_update_impl(
+        *args
+    )
+    fake_output = rope_kvcache_fusion.fused_rope_and_unified_kv_cache_update_fake(*args)
+    assert runtime_output.shape == fake_output.shape
+    assert runtime_output.dtype == fake_output.dtype
+    assert runtime_output.device == fake_output.device
+
+
 def test_rope_kvcache_fusion_default_keeps_large_ranges_unfused():
     vllm_config = VllmConfig(
         compilation_config=CompilationConfig(

vllm/compilation/passes/fusion/rope_kvcache_fusion.py

@@ -74,7 +74,7 @@ def fused_rope_and_unified_kv_cache_update_impl(
             layer_slot_mapping,
         )
 
-    return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype)
+    return query.new_empty(0)
 
 
 def fused_rope_and_unified_kv_cache_update_fake(

vllm/model_executor/layers/attention/attention.py

@@ -794,7 +794,7 @@ def unified_kv_cache_update(
             layer_slot_mapping,
         )
 
-    return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype)
+    return key.new_empty(0)
 
 
 def unified_kv_cache_update_fake(