Fix missing data type in flashinfer prefill

vllm-project · Nov 8, 2024 · ea73f6c · ea73f6c
1 parent 201fc07
commit ea73f6c
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/vllm/attention/backends/flashinfer.py b/vllm/attention/backends/flashinfer.py
@@ -364,7 +364,11 @@ def begin_forward(self):
                     self.paged_kv_indices,
                     self.paged_kv_last_page_len[:self.num_prefills],
                     self.num_qo_heads, self.num_kv_heads, self.head_dim,
-                    self.page_size)
+                    self.page_size,
+                    # kv-cache data type.
+                    data_type=self.data_type,
+                    # query data type.
+                    q_data_type=self.q_data_type)
         if self.num_decode_tokens > 0:
             assert self.paged_kv_indices is not None
             assert self.paged_kv_indptr is not None