[Core] Refactor Attention Take 2 (#3462)

2024-03-24 21:39:33 -07:00
parent b0dfa91dd7
commit 925f3332ca
47 changed files with 1268 additions and 1117 deletions
--- a/tests/kernels/test_prefix_prefill.py
+++ b/tests/kernels/test_prefix_prefill.py
@@ -3,8 +3,7 @@ import pytest
 import time

 import torch
-from vllm.model_executor.layers.attention.ops.prefix_prefill import (
-    context_attention_fwd)
+from vllm.attention.ops.prefix_prefill import context_attention_fwd
 from xformers import ops as xops
 from xformers.ops.fmha.attn_bias import BlockDiagonalCausalFromBottomRightMask