Upgrade to vllm 0.17.0 corex v4.1 overlay

vllm/v1/attention/ops/flashmla.py

@@ -6,6 +6,7 @@ import torch
 
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
+from vllm import _custom_ops as ops
 
 logger = init_logger(__name__)
 
@@ -151,7 +152,34 @@ def flash_mla_with_kvcache_fp8(
         descale_k,
     )
     return out, softmax_lse
+def flash_mla_sparse_prefill(
+    q: torch.Tensor,
+    kv: torch.Tensor,
+    indices: torch.Tensor,
+    sm_scale: float,
+    d_v: int = 512,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Sparse attention prefill kernel
 
+    Args:
+    - q: [s_q, h_q, d_qk], bfloat16
+    - kv: [s_kv, h_kv, d_qk], bfloat16
+    - indices: [s_q, h_kv, topk], int32.
+        Invalid indices should be set to -1 or numbers >= s_kv
+    - sm_scale: float
+    - d_v: The dimension of value vectors. Can only be 512
+
+    Returns:
+    - (output, max_logits, lse)
+        About the definition of output,
+        max_logits and lse, please refer to README.md
+    - output: [s_q, h_q, d_v], bfloat16
+    - max_logits:  [s_q, h_q], float
+    - lse: [s_q, h_q], float, 2-based log-sum-exp
+    """
+    results = ops.sparse_prefill_fwd(q, kv, indices,sm_scale, d_v)
+    return results
 
 #
 # TODO: Add fake functions