[compile] Fix fake kernel return dtype (#50444)
Signed-off-by: Richard Zou <zou3519@gmail.com> Co-authored-by: Codex <noreply@openai.com>
Richard Zou committed
Jul 30, 2026 at 15:50 UTC
c27b080d262ee5966cf8d362d3a9d0cefd930908
3 files changed
+34
-2
tests/compile/passes/test_rope_kvcache_fusion.py
+32
@@ -29,6 +29,7 @@ from vllm.config import (
29
from vllm.config.utils import Range
30
from vllm.forward_context import get_forward_context, set_forward_context
31
from vllm.model_executor.layers.attention import Attention
32
+from vllm.model_executor.layers.attention import attention as attention_module
33
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
34
from vllm.platforms import current_platform
35
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
44
FP8_DTYPE = current_platform.fp8_dtype()
45
46
47
+def test_kv_cache_update_ops_fake_tensor_metadata(monkeypatch: pytest.MonkeyPatch):
48
+ query = torch.empty(1, dtype=torch.bfloat16)
49
+ kv_cache = torch.empty(1, dtype=torch.uint8)
50
+ context = (None, None, kv_cache, None)
51
+ monkeypatch.setattr(attention_module, "get_attention_context", lambda _: context)
52
+ monkeypatch.setattr(rope_kvcache_fusion, "get_attention_context", lambda _: context)
53
+
54
+ runtime_output = attention_module.unified_kv_cache_update(query, query, "layer")
55
+ fake_output = attention_module.unified_kv_cache_update_fake(query, query, "layer")
56
+ assert runtime_output.shape == fake_output.shape
57
+ assert runtime_output.dtype == fake_output.dtype
58
+ assert runtime_output.device == fake_output.device
59
+
60
+ args = (
61
+ query,
62
+ query,
63
+ query,
64
+ torch.empty(1, dtype=torch.int64),
65
+ torch.empty(1),
66
+ True,
67
+ "layer",
68
+ )
69
+ runtime_output = rope_kvcache_fusion.fused_rope_and_unified_kv_cache_update_impl(
70
+ *args
71
+ )
72
+ fake_output = rope_kvcache_fusion.fused_rope_and_unified_kv_cache_update_fake(*args)
73
+ assert runtime_output.shape == fake_output.shape
74
+ assert runtime_output.dtype == fake_output.dtype
75
+ assert runtime_output.device == fake_output.device
76
+
77
+
78
def test_rope_kvcache_fusion_default_keeps_large_ranges_unfused():
79
vllm_config = VllmConfig(
80
compilation_config=CompilationConfig(
vllm/compilation/passes/fusion/rope_kvcache_fusion.py
+1
-1
@@ -74,7 +74,7 @@ def fused_rope_and_unified_kv_cache_update_impl(
74
layer_slot_mapping,
75
)
76
77
- return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype)
77
+ return query.new_empty(0)
78
79
80
def fused_rope_and_unified_kv_cache_update_fake(
vllm/model_executor/layers/attention/attention.py
+1
-1
@@ -794,7 +794,7 @@ def unified_kv_cache_update(
794
layer_slot_mapping,
795
)
796
797
- return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype)
797
+ return key.new_empty(0)
798
799
800
def unified_kv_cache_update_fake(