[wip]

csrc/cache.h

@@ -24,6 +24,14 @@ void reshape_and_cache(
   const std::string& kv_cache_dtype,
   const float kv_scale);
 
+void reshape_and_cache_flash(
+  torch::Tensor& key,
+  torch::Tensor& value,
+  torch::Tensor& key_cache,
+  torch::Tensor& value_cache,
+  torch::Tensor& slot_mapping,
+  const std::string& kv_cache_dtype);
+
 // Just for unittest
 void convert_fp8(
   torch::Tensor& src_cache,

csrc/cache_kernels.cu

@@ -215,6 +215,42 @@ __global__ void reshape_and_cache_kernel(
   }
 }
 
+template<typename scalar_t>
+__global__ void reshape_and_cache_flash_kernel(
+  const scalar_t* __restrict__ key,           // [num_tokens, num_heads, head_size]
+  const scalar_t* __restrict__ value,         // [num_tokens, num_heads, head_size]
+  scalar_t* __restrict__ k_cache,             // [num_blocks, block_size, num_heads, head_size]
+  scalar_t* __restrict__ v_cache,             // [num_blocks, block_size, num_heads, head_size]
+  const int64_t* __restrict__ slot_mapping,   // [num_tokens]
+  const int block_stride,
+  const int key_stride,
+  const int value_stride,
+  const int num_heads,
+  const int head_size,
+  const int block_size) {
+  const int64_t token_idx = blockIdx.x;
+  const int64_t slot_idx = slot_mapping[token_idx];
+  // NOTE: slot_idx can be -1 if the token is padded
+  if (slot_idx < 0) {
+    return;
+  }
+  const int64_t block_idx = slot_idx / block_size;
+  const int64_t block_offset = slot_idx % block_size;
+  const int n = num_heads * head_size;
+  for (int i = threadIdx.x; i < n; i += blockDim.x) {
+    const int64_t src_key_idx = token_idx * key_stride + i;
+    const int64_t src_value_idx = token_idx * value_stride + i;
+    const int head_idx = i / head_size;
+    const int head_offset = i % head_size;
+    const int64_t tgt_value_idx = block_idx * block_stride
+                              + block_offset * num_heads * head_size
+                              + head_idx * head_size
+                              + head_offset;
+    k_cache[tgt_value_idx] = key[src_key_idx];
+    v_cache[tgt_value_idx] = value[src_value_idx];
+  }
+}
+
 } // namespace vllm
 
 #define CALL_RESHAPE_AND_CACHE(KV_T, CACHE_T, IS_FP8_KV_CACHE)                                     \
@@ -275,6 +311,50 @@ void reshape_and_cache(
   }
 }
 
+void reshape_and_cache_flash(
+  torch::Tensor& key,           // [num_tokens, num_heads, head_size]
+  torch::Tensor& value,         // [num_tokens, num_heads, head_size]
+  torch::Tensor& k_cache,       // [num_blocks, block_size, num_heads, head_size]
+  torch::Tensor& v_cache,       // [num_blocks, block_size, num_heads, head_size]
+  torch::Tensor& slot_mapping,  // [num_tokens]
+  const std::string& kv_cache_dtype)
+{
+  if (kv_cache_dtype != "auto") {
+    TORCH_CHECK(false, "Unsupported data type of kv cache: ", kv_cache_dtype);
+  }
+  int num_tokens = key.size(0);
+  int num_heads = key.size(1);
+  int head_size = key.size(2);
+  int block_size = k_cache.size(1);
+
+  int key_stride = key.stride(0);
+  int value_stride = value.stride(0);
+  int block_stride = k_cache.stride(0);
+  TORCH_CHECK(k_cache.stride(0) == v_cache.stride(0));
+
+  dim3 grid(num_tokens);
+  dim3 block(std::min(num_heads * head_size, 512));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  VLLM_DISPATCH_FLOATING_TYPES(
+    key.scalar_type(),
+    "reshape_and_cache_flash",
+    [&] {
+      vllm::reshape_and_cache_flash_kernel<scalar_t><<<grid, block, 0, stream>>>(
+        key.data_ptr<scalar_t>(),
+        value.data_ptr<scalar_t>(),
+        k_cache.data_ptr<scalar_t>(),
+        v_cache.data_ptr<scalar_t>(),
+        slot_mapping.data_ptr<int64_t>(),
+        block_stride,
+        key_stride,
+        value_stride,
+        num_heads,
+        head_size,
+        block_size);
+    });
+}
+
 namespace vllm {
 
 template<typename Tout, typename Tin>

csrc/pybind.cpp

@@ -96,6 +96,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     "reshape_and_cache",
     &reshape_and_cache,
     "Reshape the key and value tensors and cache them");
+  cache_ops.def(
+    "reshape_and_cache_flash",
+    &reshape_and_cache_flash,
+    "Reshape the key and value tensors and cache them");
   cache_ops.def(
     "convert_fp8",
     &convert_fp8,

requirements-cuda.txt

@@ -5,5 +5,6 @@
 ray >= 2.9
 nvidia-ml-py # for pynvml package
 vllm-nccl-cu12>=2.18,<2.19  # for downloading nccl library
-torch == 2.3.0
-xformers == 0.0.26.post1  # Requires PyTorch 2.3.0
+torch == 2.2.1
+xformers == 0.0.25  # Requires PyTorch 2.2.1
+vllm-flash-attn == 2.5.7  # Requires PyTorch 2.2.1

setup.py

@@ -342,14 +342,18 @@ def _read_requirements(filename: str) -> List[str]:
 
     if _is_cuda():
         requirements = _read_requirements("requirements-cuda.txt")
-        cuda_major = torch.version.cuda.split(".")[0]
+        cuda_major, cuda_minor = torch.version.cuda.split(".")
         modified_requirements = []
         for req in requirements:
             if "vllm-nccl-cu12" in req:
-                modified_requirements.append(
-                    req.replace("vllm-nccl-cu12", f"vllm-nccl-cu{cuda_major}"))
-            else:
-                modified_requirements.append(req)
+                req = req.replace("vllm-nccl-cu12",
+                                  f"vllm-nccl-cu{cuda_major}")
+            elif ("vllm-flash-attn" in req
+                  and not (cuda_major == "12" and cuda_minor == "1")):
+                # vllm-flash-attn is built only for CUDA 12.1.
+                # Skip for other versions.
+                continue
+            modified_requirements.append(req)
         requirements = modified_requirements
     elif _is_hip():
         requirements = _read_requirements("requirements-rocm.txt")

tests/kernels/test_cache.py

@@ -14,6 +14,7 @@
 NUM_HEADS = [8]  # Arbitrary values for testing
 HEAD_SIZES = [64, 80, 96, 112, 128, 256]
 BLOCK_SIZES = [8, 16, 32]
+PADDINGS = [8, 16, 0]
 
 # Arbitrary values for testing
 # don't make it too large. e.g. [1024, 36000] will OOM
@@ -296,3 +297,76 @@ def test_fp8_conversion(
     ops.convert_fp8(cache_fp8, converted_cache)
 
     assert torch.allclose(cache, converted_cache, atol=0.001, rtol=0.1)
+
+
+# @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
+# @pytest.mark.parametrize("num_heads", NUM_HEADS)
+# @pytest.mark.parametrize("head_size", HEAD_SIZES)
+# @pytest.mark.parametrize("block_size", BLOCK_SIZES)
+# @pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
+# @pytest.mark.parametrize("dtype", DTYPES)
+# @pytest.mark.parametrize("seed", SEEDS)
+# @pytest.mark.parametrize("padding", PADDINGS)
+# @torch.inference_mode()
+# def test_reshape_and_cache_flash(
+#     kv_cache_factory,
+#     num_tokens: int,
+#     num_heads: int,
+#     head_size: int,
+#     block_size: int,
+#     num_blocks: int,
+#     dtype: torch.dtype,
+#     seed: int,
+#     padding: int,
+# ) -> None:
+#     random.seed(seed)
+#     torch.random.manual_seed(seed)
+#     torch.cuda.manual_seed(seed)
+
+#     # Create a random slot mapping.
+#     num_slots = block_size * num_blocks
+#     slot_mapping = random.sample(range(num_slots), num_tokens)
+#     slot_mapping = torch.tensor(slot_mapping, dtype=torch.long, device='cuda')
+
+#     qkv = torch.randn(num_tokens,
+#                       3,
+#                       num_heads,
+#                       head_size,
+#                       dtype=dtype,
+#                       device='cuda')
+#     _, key, value = qkv.unbind(dim=1)
+
+#     # Create the KV caches.
+#     key_caches, value_caches = kv_cache_factory(num_blocks,
+#                                                 block_size,
+#                                                 1,
+#                                                 num_heads,
+#                                                 head_size,
+#                                                 dtype,
+#                                                 seed,
+#                                                 flash_style=True)
+#     key_cache, value_cache = key_caches[0], value_caches[0]
+
+#     # Clone the KV caches.
+#     cloned_key_cache = key_cache.clone()
+#     cloned_value_cache = value_cache.clone()
+#     num_tokens = torch.tensor([num_tokens], dtype=torch.long, device="cuda")
+
+#     padded_key, padded_value = pad_key_value(key, value, padding)
+#     # Call the reshape_and_cache kernel.
+#     cache_ops.reshape_and_cache_flash(padded_key, padded_value, key_cache,
+#                                       value_cache, slot_mapping, num_tokens)
+
+#     # Run the reference implementation.
+#     block_indicies = torch.div(slot_mapping, block_size, rounding_mode='floor')
+#     block_indicies = block_indicies.cpu().tolist()
+#     block_offsets = slot_mapping % block_size
+#     block_offsets = block_offsets.cpu().tolist()
+#     for i in range(num_tokens):
+#         block_idx = block_indicies[i]
+#         block_offset = block_offsets[i]
+#         cloned_key_cache[block_idx, block_offset, :, :] = key[i]
+#         cloned_value_cache[block_idx, block_offset, :, :] = value[i]
+
+#     assert torch.allclose(key_cache, cloned_key_cache)
+#     assert torch.allclose(value_cache, cloned_value_cache)

tests/kernels/test_flash_attn.py

@@ -0,0 +1,122 @@
+import random
+
+import pytest
+import torch
+from vllm_flash_attn import flash_attn_varlen_func
+from xformers import ops as xops
+from xformers.ops.fmha.attn_bias import BlockDiagonalCausalFromBottomRightMask
+
+NUM_HEADS = [8]
+NUM_QUERIES_PER_KV = [1]
+HEAD_SIZES = [128]
+DTYPES = [torch.float16]
+
+
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
+@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
+@pytest.mark.parametrize("dtype", DTYPES)
+@torch.inference_mode()
+def test_flashinfer_append(num_heads: int, num_queries_per_kv: int,
+                           head_size: int, dtype: torch.dtype):
+    random.seed(0)
+    torch.manual_seed(0)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(0)
+    torch.set_default_device("cuda")
+    batch_size = 10
+    cache_size = 640
+    block_size = 16
+    prefix_lens = [random.randint(100, 200) for _ in range(batch_size)]
+    append_lens = [random.randint(1, 5) for _ in range(batch_size)]
+    seq_lens = [a + b for a, b in zip(prefix_lens, append_lens)]
+    num_tokens = sum(append_lens)
+    query = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
+    query.uniform_(-1e-1, 1e-1)
+    num_kv_heads = num_heads // num_queries_per_kv
+    key_value = torch.empty(sum(seq_lens),
+                            2,
+                            num_kv_heads,
+                            head_size,
+                            dtype=dtype)
+    key_value.uniform_(-1e-1, 1e-1)
+    key, value = key_value.unbind(dim=1)
+    values = torch.arange(0, cache_size, dtype=torch.int32)
+    values = values[torch.randperm(cache_size)]
+    max_block_per_request = int(cache_size / batch_size)
+    block_table = values[:batch_size * max_block_per_request].view(
+        batch_size, max_block_per_request)
+    k_cache = torch.zeros(cache_size,
+                          block_size,
+                          num_kv_heads,
+                          head_size,
+                          dtype=dtype)
+    v_cache = torch.zeros(cache_size,
+                          block_size,
+                          num_kv_heads,
+                          head_size,
+                          dtype=dtype)
+    qo_indptr = torch.cumsum(torch.tensor([0] + append_lens),
+                             dim=0,
+                             dtype=torch.int32)
+    seq_start_loc = torch.cumsum(torch.tensor([0] + seq_lens),
+                                 dim=0,
+                                 dtype=torch.int32)
+    paged_kv_last_page_len = []
+    paged_kv_indptr = [0]
+    page_kv_indices = []
+    total_block_num = 0
+    for i in range(batch_size):
+        # copy key, value to kv cache
+        cur_prefix_id = 0
+        block_id = 0
+        while cur_prefix_id < seq_lens[i]:
+            start_loc = seq_start_loc[i] + cur_prefix_id
+            if cur_prefix_id + block_size > seq_lens[i]:
+                end_loc = seq_start_loc[i] + seq_lens[i]
+            else:
+                end_loc = start_loc + block_size
+            start_slot = block_table[i, block_id] * block_size
+            end_slot = start_slot + end_loc - start_loc
+            k_cache.view(-1, num_kv_heads,
+                         head_size)[start_slot:end_slot].copy_(
+                             key[start_loc:end_loc])
+            v_cache.view(-1, num_kv_heads,
+                         head_size)[start_slot:end_slot].copy_(
+                             value[start_loc:end_loc])
+            cur_prefix_id += block_size
+            block_id += 1
+        paged_kv_last_page_len.append((seq_lens[i] - 1) % block_size + 1)
+        cur_block_num = (seq_lens[i] - 1) // block_size + 1
+        page_kv_indices.extend(block_table[i, :cur_block_num])
+        total_block_num += cur_block_num
+        paged_kv_indptr.append(total_block_num)
+    output = flash_attn_varlen_func(
+        query,
+        k_cache,
+        v_cache,
+        cu_seqlens_q=qo_indptr,
+        cu_seqlens_k=seq_start_loc,
+        max_seqlen_q=max(append_lens),
+        max_seqlen_k=max(seq_lens),
+        causal=True,
+        block_table=block_table,
+    )
+    query = query.unsqueeze(0)
+    key = key.unsqueeze(0)
+    value = value.unsqueeze(0)
+    attn_bias = BlockDiagonalCausalFromBottomRightMask.from_seqlens(
+        append_lens, seq_lens)
+    scale = float(1.0 / (head_size**0.5))
+    attn_op = xops.fmha.cutlass.FwOp()
+    output_ref = xops.memory_efficient_attention_forward(
+        query,
+        key,
+        value,
+        attn_bias=attn_bias,
+        p=0.0,
+        scale=scale,
+        op=attn_op,
+    ).squeeze(0)
+    print((output - output_ref).abs().max())
+    assert torch.allclose(output_ref, output, atol=1e-4, rtol=1e-2)