[Kernel][Backend][Model] Blocksparse flash attention kernel and Phi-3-Small model

csrc/ops.h

@@ -15,7 +15,12 @@ void paged_attention_v1(
   int max_seq_len,
   const c10::optional<torch::Tensor>& alibi_slopes,
   const std::string& kv_cache_dtype,
-  float kv_scale);
+  float kv_scale,
+  const int tp_rank,
+  const int blocksparse_local_blocks,
+  const int blocksparse_vert_stride,
+  const int blocksparse_block_size,
+  const int blocksparse_head_sliding_step);
 
 void paged_attention_v2(
   torch::Tensor& out,
@@ -33,7 +38,12 @@ void paged_attention_v2(
   int max_seq_len,
   const c10::optional<torch::Tensor>& alibi_slopes,
   const std::string& kv_cache_dtype,
-  float kv_scale);
+  float kv_scale,
+  const int tp_rank,
+  const int blocksparse_local_blocks,
+  const int blocksparse_vert_stride,
+  const int blocksparse_block_size,
+  const int blocksparse_head_sliding_step);
 
 void rms_norm(
   torch::Tensor& out,

docs/source/models/supported_models.rst

@@ -123,6 +123,10 @@ Alongside each architecture, we include some popular models that use it.
     - Phi-3
     - :code:`microsoft/Phi-3-mini-4k-instruct`, :code:`microsoft/Phi-3-mini-128k-instruct`, etc.
     -
+  * - :code:`Phi3SmallForCausalLM`
+    - Phi-3-Small
+    - :code:`microsoft/Phi-3-small-8k-instruct`, :code:`microsoft/Phi-3-small-128k-instruct`, etc.
+    -
   * - :code:`QWenLMHeadModel`
     - Qwen
     - :code:`Qwen/Qwen-7B`, :code:`Qwen/Qwen-7B-Chat`, etc.

vllm/_custom_ops.py

@@ -45,11 +45,17 @@ def paged_attention_v1(
     alibi_slopes: Optional[torch.Tensor],
     kv_cache_dtype: str,
     kv_scale: float,
+    tp_rank: int,
+    blocksparse_local_blocks: int,
+    blocksparse_vert_stride: int,
+    blocksparse_block_size: int,
+    blocksparse_head_sliding_step: int,
 ) -> None:
-    vllm_ops.paged_attention_v1(out, query, key_cache, value_cache,
-                                num_kv_heads, scale, block_tables, seq_lens,
-                                block_size, max_seq_len, alibi_slopes,
-                                kv_cache_dtype, kv_scale)
+    vllm_ops.paged_attention_v1(
+        out, query, key_cache, value_cache, num_kv_heads, scale, block_tables,
+        seq_lens, block_size, max_seq_len, alibi_slopes, kv_cache_dtype,
+        kv_scale, tp_rank, blocksparse_local_blocks, blocksparse_vert_stride,
+        blocksparse_block_size, blocksparse_head_sliding_step)
 
 
 def paged_attention_v2(
@@ -69,12 +75,18 @@ def paged_attention_v2(
     alibi_slopes: Optional[torch.Tensor],
     kv_cache_dtype: str,
     kv_scale: float,
+    tp_rank: int,
+    blocksparse_local_blocks: int,
+    blocksparse_vert_stride: int,
+    blocksparse_block_size: int,
+    blocksparse_head_sliding_step: int,
 ) -> None:
-    vllm_ops.paged_attention_v2(out, exp_sum, max_logits, tmp_out, query,
-                                key_cache, value_cache, num_kv_heads, scale,
-                                block_tables, seq_lens, block_size,
-                                max_seq_len, alibi_slopes, kv_cache_dtype,
-                                kv_scale)
+    vllm_ops.paged_attention_v2(
+        out, exp_sum, max_logits, tmp_out, query, key_cache, value_cache,
+        num_kv_heads, scale, block_tables, seq_lens, block_size, max_seq_len,
+        alibi_slopes, kv_cache_dtype, kv_scale, tp_rank,
+        blocksparse_local_blocks, blocksparse_vert_stride,
+        blocksparse_block_size, blocksparse_head_sliding_step)
 
 
 # pos encoding ops

vllm/attention/backends/abstract.py

@@ -114,6 +114,7 @@ def __init__(
         num_kv_heads: Optional[int] = None,
         alibi_slopes: Optional[List[float]] = None,
         sliding_window: Optional[int] = None,
+        blocksparse_params: Optional[Dict[str, Any]] = None,
         kv_cache_dtype: str = "auto",
     ) -> None:
         raise NotImplementedError