Repository: bytedance-iaas/vllm
Pull request #99: feat: add sparse-prefill
Author: login: artetaout
Created: 2025-04-16T06:33:27Z | Merged: 2025-04-21T16:22:04Z



<!--- pyml disable-next-line no-emphasis-as-heading -->



Diff:
diff --git a/benchmarks/benchmark_latency.py b/benchmarks/benchmark_latency.py
index dfd9bb1e6a4d..1abb5e9c99e7 100644
--- a/benchmarks/benchmark_latency.py
+++ b/benchmarks/benchmark_latency.py
@@ -14,6 +14,7 @@
 from benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json
 from tqdm import tqdm
 
+import vllm
 from vllm import LLM, SamplingParams
 from vllm.engine.arg_utils import EngineArgs
 from vllm.inputs import PromptType
@@ -37,7 +38,9 @@ def main(args: argparse.Namespace):
     print(args)
 
     engine_args = EngineArgs.from_cli_args(args)
-
+    if vllm.envs.VLLM_USE_SP_PREFILL:
+        engine_args.enable_chunked_prefill = False
+        engine_args.max_num_seqs = 1
     # NOTE(woosuk): If the request cannot be processed in a single batch,
     # the engine will automatically process the request in multiple batches.
     llm = LLM(**dataclasses.asdict(engine_args))
diff --git a/vllm/attention/backends/flash_attn.py b/vllm/attention/backends/flash_attn.py
index f9c5ad4df54e..1283c7bcc4e4 100755
--- a/vllm/attention/backends/flash_attn.py
+++ b/vllm/attention/backends/flash_attn.py
@@ -34,6 +34,9 @@
     from vllm.worker.model_runner import (ModelInputForGPUBuilder,
                                           ModelInputForGPUWithSamplingMetadata)
 
+import vllm.attention.backends.sparse_prefill as sparse_prefill
+import vllm.envs
+
 logger = init_logger(__name__)
 
 
@@ -779,6 +782,15 @@ def forward(
             # Prompt run.
             if (kv_cache.numel() == 0 or prefill_meta.block_tables is None
                     or prefill_meta.block_tables.numel() == 0):
+
+                if vllm.envs.VLLM_USE_SP_PREFILL:
+                    out = sparse_prefill.sparse_prefill_attention(
+                        query.unsqueeze(0), key.unsqueeze(0),
+                        value.unsqueeze(0)).squeeze(0)
+                    assert output[:attn_metadata.
+                                  num_prefill_tokens].shape == out.shape
+                    output[:attn_metadata.num_prefill_tokens] = out
+                    return output
                 # normal attention
                 # When block_tables are not filled, it means q and k are the
                 # prompt, and they have the same length.
diff --git a/vllm/attention/backends/sparse_prefill.py b/vllm/attention/backends/sparse_prefill.py
new file mode 100644
index 000000000000..bc8d7662fb5b
--- /dev/null
+++ b/vllm/attention/backends/sparse_prefill.py
@@ -0,0 +1,768 @@
+# SPDX-License-Identifier: Apache-2.0
+import math
+from typing import List, Optional, Union
+
+import torch
+import triton
+import triton.language as tl
+from einops import rearrange
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+try:
+    from flash_attn import flash_attn_func
+except ImportError as e:
+    logger.warning('Import error %s, please install flash-attn correctly',
+                   e.msg)
+
+causal_mask = None
+
+
+def gpu_info():
+    if torch.cuda.is_available():
+        device_name = torch.cuda.get_device_name(0).lower()
+        device_capability = torch.cuda.get_device_capability()
+        major, minor = device_capability
+        return device_name, major
+    return None, None
+
+
+GPU_NAME, GPU_MAJOR = gpu_info()
+
+
+def get_num_warps_stages(head_dim, block_size, gpu_name):
+    gpu_name = gpu_name.lower()
+    # Determine if head_dim and block_size exceed 64
+    head_large = head_dim > 64
+    block_large = block_size > 64
+
+    if "h100" in gpu_name:
+        # Hopper GPU recommendations
+        if head_large and block_large:
+            num_warps = 8
+            num_stages = 3
+        elif head_large or block_large:
+            num_warps = 4
+            num_stages = 3
+        else:
+            num_warps = 2
+            num_stages = 2
+    elif "a100" in gpu_name:
+        # Ampere GPU recommendations
+        if head_large and block_large or head_large or block_large:
+            num_warps = 8
+            num_stages = 3
+        else:
+            num_warps = 2
+            num_stages = 2
+    elif "4090" in gpu_name:
+        if head_large and block_large:
+            num_warps = 8
+            num_stages = 2
+        elif head_large or block_large:
+            num_warps = 8
+            num_stages = 3
+        else:
+            num_warps = 2
+            num_stages = 2
+    else:
+        # use default setting, maybe not optimal
+        if head_large and block_large:
+            num_warps = 8
+            num_stages = 2
+        elif head_large or block_large:
+            num_warps = 4
+            num_stages = 3
+        else:
+            num_warps = 2
+            num_stages = 2
+    return num_warps, num_stages
+
+
+@triton.jit
+def block_wise_prefill_attention_kernel(
+        q_ptr,  # shape: [batch_size, seq_len, num_heads, head_dim]
+        k_ptr,
+        v_ptr,
+        o_ptr,
+        block_idx_ptr,  # shape: [batch_size, num_heads, num_all_block]
+        idx_bin_ptr,  # shape: [batch_size, num_heads, seq_len / block_size + 1]
+        # shape
+    BATCH_SIZE,
+        NUM_HEADS,
+        NUM_KV_HEADS,
+        NUM_SHARE_Q_HEADS,
+        Q_LEN,
+        K_LEN,
+        HEAD_DIM: tl.constexpr,
+        NUM_BLOCK,
+        grid_offset,
+        # softmax_scale
+        softmax_scale,
+        # gqa
+        gqa_interleave: tl.constexpr,
+        # stride
+        stride_qb,
+        stride_qn,
+        stride_qh,
+        stride_qd,
+        stride_kb,
+        stride_kn,
+        stride_kh,
+        stride_kd,
+        stride_vb,
+        stride_vn,
+        stride_vh,
+        stride_vd,
+        stride_ob,
+        stride_on,
+        stride_oh,
+        stride_od,
+        stride_bb,
+        stride_bh,
+        stride_bt,
+        stride_ib,
+        stride_ih,
+        stride_it,
+        # META parameters
+        BLOCK_SIZE_Q: tl.constexpr,  # q block size
+        BLOCK_SIZE_K: tl.constexpr,  # k block size
+):
+    tl.static_assert(BLOCK_SIZE_Q == BLOCK_SIZE_K)
+    pid_b = tl.program_id(0)
+    pid_h = tl.program_id(1)
+    if gqa_interleave:
+        pid_kh = pid_h % NUM_KV_HEADS
+    else:
+        pid_kh = pid_h // NUM_SHARE_Q_HEADS
+    pid_q = tl.program_id(2)
+    # get column index bin
+    idx_bin_ptr = idx_bin_ptr + pid_b * stride_ib + pid_h * stride_ih
+    bin_start = tl.load(idx_bin_ptr + pid_q * stride_it)
+    bin_end = tl.load(idx_bin_ptr + (pid_q + 1) * stride_it)
+    num_active_block = bin_end - bin_start
+    # get column block index ptr
+    block_idx_ptr = (block_idx_ptr + pid_b * stride_bb + pid_h * stride_bh +
+                     bin_start * stride_bt)
+    # init qkv ptrs
+    q_ptrs = tl.make_block_ptr(
+        base=q_ptr + pid_b * stride_qb + pid_h * stride_qh,
+        shape=(Q_LEN, HEAD_DIM),
+        strides=(stride_qn, stride_qd),
+        offsets=(pid_q * BLOCK_SIZE_Q - grid_offset, 0),
+        block_shape=(BLOCK_SIZE_Q, HEAD_DIM),
+        order=(1, 0),
+    )
+    k_ptrs = tl.make_block_ptr(
+        base=k_ptr + pid_b * stride_kb + pid_kh * stride_kh,
+        shape=(HEAD_DIM, K_LEN),
+        strides=(stride_kd, stride_kn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_SIZE_K),
+        order=(0, 1),
+    )
+    v_ptrs = tl.make_block_ptr(
+        base=v_ptr + pid_b * stride_vb + pid_kh * stride_vh,
+        shape=(K_LEN, HEAD_DIM),
+        strides=(stride_vn, stride_vd),
+        offsets=(0, 0),
+        block_shape=(BLOCK_SIZE_K, HEAD_DIM),
+        order=(1, 0),
+    )
+    # load q
+    q = tl.load(q_ptrs, boundary_check=(0, 1), padding_option="zero")
+    # init statistics
+    off_m = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q - grid_offset
+    off_n = tl.arange(0, BLOCK_SIZE_K)
+    m_i = tl.full((BLOCK_SIZE_Q, ), float("-inf"), dtype=tl.float32)
+    lse_i = tl.full((BLOCK_SIZE_Q, ), float("-inf"), dtype=tl.float32)
+    acc_o = tl.full((BLOCK_SIZE_Q, HEAD_DIM), 0, dtype=tl.float32)
+    # flash attention
+    for i in range(0, num_active_block):
+        # get current block start index
+        c = tl.load(block_idx_ptr).to(
+            tl.int32) % NUM_BLOCK * BLOCK_SIZE_K - grid_offset
+        block_idx_ptr = block_idx_ptr + stride_bt
+        # load k
+        k = tl.load(tl.advance(k_ptrs, (0, c)),
+                    boundary_check=(1, ),
+                    padding_option="zero")
+        # compute qk
+        qk = tl.zeros((BLOCK_SIZE_Q, BLOCK_SIZE_K), dtype=tl.float32)
+        qk += tl.where((c + off_n)[None, :] >= 0, 0, float("-inf"))
+        qk += tl.where(off_m[:, None] >= (c + off_n)[None, :], 0,
+                       float("-inf"))
+        qk += tl.dot(q, k) * softmax_scale
+        # compute m_ij and l_ij
+        m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
+        p = tl.math.exp2(qk - m_ij[:, None])
+        l_ij = tl.sum(p, axis=1)
+        # scale acc_o
+        acc_o_scale = tl.math.exp2(m_i - m_ij)
+        acc_o = acc_o * acc_o_scale[:, None]
+        # load v and update acc_o
+        v = tl.load(tl.advance(v_ptrs, (c, 0)),
+                    boundary_check=(0, ),
+                    padding_option="zero")
+        p = p.to(v.dtype)
+        acc_o += tl.dot(p, v)
+        # update statistics
+        m_i = m_ij
+        lse_i = m_ij + tl.math.log2(tl.math.exp2(lse_i - m_ij) + l_ij)
+    # final scale
+    acc_o = acc_o * tl.math.exp2(m_i - lse_i)[:, None]
+    # save output
+    o_ptrs = tl.make_block_ptr(
+        base=o_ptr + pid_b * stride_ob + pid_h * stride_oh,
+        shape=(Q_LEN, HEAD_DIM),
+        strides=(stride_on, stride_od),
+        offsets=(pid_q * BLOCK_SIZE_Q - grid_offset, 0),
+        block_shape=(BLOCK_SIZE_Q, HEAD_DIM),
+        order=(1, 0),
+    )
+    tl.store(o_ptrs, acc_o.to(tl.bfloat16), boundary_check=(0, ))
+
+
+@triton.jit
+def count_kernel(
+    x_ptr,
+    y_ptr,
+    k,
+    r,
+    stride_xb,
+    stride_xh,
+    stride_xk,
+    stride_yb,
+    stride_yh,
+    stride_yr,
+    BLOCK_SIZE_K: tl.constexpr,
+    BLOCK_SIZE_R: tl.constexpr,
+):
+    pid_b = tl.program_id(0)
+    pid_h = tl.program_id(1)
+    x_ptr = x_ptr + pid_b * stride_xb + pid_h * stride_xh
+    off_k = tl.arange(0, BLOCK_SIZE_K)
+    x_ptrs = x_ptr + off_k * stride_xk
+    y = tl.zeros((BLOCK_SIZE_R, ), dtype=tl.int32)
+    for i in range(0, k, BLOCK_SIZE_K):
+        x = tl.load(x_ptrs, off_k < k - i, -1)
+        x = x // r
+        x = tl.where(off_k < k - i, x, -1)
+        y += tl.histogram(x, BLOCK_SIZE_R)
+        x_ptrs = x_ptrs + BLOCK_SIZE_K * stride_xk
+    y = tl.cumsum(y, axis=0)
+    y_ptr = y_ptr + pid_b * stride_yb + pid_h * stride_yh + stride_yr
+    off_r = tl.arange(0, BLOCK_SIZE_R)
+    tl.store(y_ptr + off_r * stride_yr, y, off_r < r)
+
+
+def triton_column_count_cumsum(x: torch.Tensor,
+                               num_columns: int) -> torch.Tensor:
+    x = x.to(torch.int32)
+    b, h, k = x.shape
+    r = num_columns
+    block_size_k = min(triton.next_power_of_2(k), 4096)
+    block_size_r = triton.next_power_of_2(r + 2)
+    y = torch.zeros(b, h, r + 1, device=x.device, dtype=torch.int32)
+    count_kernel[(b, h)](
+        x,
+        y,
+        k,
+        r,
+        x.stride(0),
+        x.stride(1),
+        x.stride(2),
+        y.stride(0),
+        y.stride(1),
+        y.stride(2),
+        block_size_k,
+        block_size_r,
+    )
+    return y
+
+
+def triton_block_wise_prefill_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_idx: Union[torch.Tensor, List[List[torch.Tensor]]],
+    block_size: int,
+    grid_offset: int = 0,
+    softmax_scale: Optional[float] = None,
+    gqa_interleave: bool = False,
+) -> torch.Tensor:
+
+    batch_size, q_len, num_q_heads, head_dim = q.shape
+    batch_size, k_len, num_kv_heads, head_dim = k.shape
+    assert q.dtype == torch.bfloat16
+    assert q_len == k_len
+    assert head_dim in {16, 32, 64,
+                        128}, "only support head_dim in {16, 32, 64, 128}"
+    assert block_size in {
+        32,
+        64,
+        128,
+    }, "only support block size in {32, 64, 128}"
+    total_q_blocks = triton.cdiv(grid_offset, block_size) + triton.cdiv(
+        q_len - grid_offset, block_size)
+    total_k_blocks = triton.cdiv(grid_offset, block_size) + triton.cdiv(
+        k_len - grid_offset, block_size)
+    # pad block_idx if get list[list[tensor]]
+    if not isinstance(block_idx, torch.Tensor):
+        assert (isinstance(block_idx, list) and isinstance(block_idx[0], list)
+                and isinstance(block_idx[0][0], torch.Tensor))
+        assert len(block_idx) == batch_size and len(
+            block_idx[0]) == num_q_heads
+        block_idx = [
+            item.view(-1, 1) for sublist in block_idx for item in sublist
+        ]
+        block_idx = torch.nn.utils.rnn.pad_sequence(
+            block_idx,
+            batch_first=True,
+            padding_value=total_k_blocks * (total_k_blocks + 1),
+            # padding_value=0,
+        )
+        block_idx = block_idx.view(batch_size, num_q_heads, -1)
+    batch_size, num_q_heads, num_block = block_idx.shape
+    assert q_len == k_len
+    assert num_block <= total_q_blocks * (total_q_blocks + 1) // 2
+    # gqa
+    assert num_q_heads % num_kv_heads == 0
+    num_share_q_heads = num_q_heads // num_kv_heads
+    # softmax_scale
+    if softmax_scale is None:
+        softmax_scale = 1 / math.sqrt(head_dim) * math.log2(math.e)
+    else:
+        softmax_scale = softmax_scale * math.log2(math.e)
+    # sort idx and get block index bins
+    block_idx = block_idx.sort(-1).values
+    idx_bins = triton_column_count_cumsum(block_idx, total_k_blocks)
+    # launch attention kernel
+    o = torch.empty_like(q)
+    num_warps, num_stages = get_num_warps_stages(head_dim, block_size,
+                                                 GPU_NAME)
+    block_wise_prefill_attention_kernel[(batch_size, num_q_heads,
+                                         total_q_blocks)](
+                                             q,
+                                             k,
+                                             v,
+                                             o,
+                                             block_idx,
+                                             idx_bins,
+                                             batch_size,
+                                             num_q_heads,
+                                             num_kv_heads,
+                                             num_share_q_heads,
+                                             q_len,
+                                             k_len,
+                                             head_dim,
+                                             total_q_blocks,
+                                             grid_offset,
+                                             softmax_scale,
+                                             gqa_interleave,
+                                             q.stride(0),
+                                             q.stride(1),
+                                             q.stride(2),
+                                             q.stride(3),
+                                             k.stride(0),
+                                             k.stride(1),
+                                             k.stride(2),
+                                             k.stride(3),
+                                             v.stride(0),
+                                             v.stride(1),
+                                             v.stride(2),
+                                             v.stride(3),
+                                             o.stride(0),
+                                             o.stride(1),
+                                             o.stride(2),
+                                             o.stride(3),
+                                             block_idx.stride(0),
+                                             block_idx.stride(1),
+                                             block_idx.stride(2),
+                                             idx_bins.stride(0),
+                                             idx_bins.stride(1),
+                                             idx_bins.stride(2),
+                                             BLOCK_SIZE_Q=block_size,
+                                             BLOCK_SIZE_K=block_size,
+                                             num_warps=num_warps,
+                                             num_stages=num_stages,
+                                         )
+    return o
+
+
+def triton_block_wise_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_idx: torch.Tensor,
+    block_size: int,
+    grid_offset: int = 0,
+    softmax_scale: Optional[float] = None,
+    gqa_interleave: bool = False,
+) -> torch.Tensor:
+
+    if q.shape[1] > 1:
+        return triton_block_wise_prefill_attention(
+            q,
+            k,
+            v,
+            block_idx,
+            block_size,
+            grid_offset,
+            softmax_scale,
+            gqa_interleave,
+        )
+
+
+def score_cover_idx(x: torch.Tensor, score: float, padding_value=0):
+    x, idx = torch.sort(x, dim=-1, descending=True)
+    cumsum_x = torch.cumsum(x, dim=-1)
+    idx[cumsum_x > score] = padding_value
+    return idx
+
+
+def transform_veritcal_slash_idx(v_idx, s_idx, num_blocks):
+    batch_size, num_heads, _ = v_idx.shape
+    range_blocks = torch.arange(num_blocks, device=s_idx.device)[None, None, :,
+                                                                 None]
+    # vertical
+    v_idx = (
+        torch.arange(0, num_blocks, device=v_idx.device)[None, None, :, None] *
+        num_blocks + v_idx[:, :, None, :]).view(batch_size, num_heads, -1)
+    v_idx[v_idx // num_blocks < v_idx % num_blocks] = 0
+    # slash
+    s_idx = (range_blocks * num_blocks + range_blocks +
+             s_idx[:, :, None, :] * num_blocks).view(batch_size, num_heads, -1)
+    s_idx[s_idx >= num_blocks * num_blocks] = 0
+    # union
+    vs_idx = torch.cat((s_idx, v_idx), dim=-1)
+    block_idx = [[torch.unique(vs_idx[b, h]) for h in range(num_heads)]
+                 for b in range(batch_size)]
+    return block_idx
+
+
+def square_root_js_divergence(p: torch.Tensor, q: torch.Tensor):
+    m = (p + q) / 2
+    return torch.sqrt(0.5 * (p * torch.log(p / m)).sum(-1) + 0.5 *
+                      (q * torch.log(q / m)).sum(-1))
+
+
+@triton.jit
+def bnhd_pool_kernel(
+        x_ptr,
+        y_ptr,
+        # pool type. avg: 0, max: 1, min: 2, max abs: 3, sum: 4
+        pool_type: tl.constexpr,
+        # shape
+        batch_size,
+        seq_len,
+        num_heads,
+        head_dim: tl.constexpr,
+        # stride
+        stride_xb,
+        stride_xn,
+        stride_xh,
+        stride_xd,
+        stride_yb,
+        stride_yn,
+        stride_yh,
+        stride_yd,
+        # META parameters
+        BLOCK_SIZE_N: tl.constexpr,
+        BLOCK_SIZE_H: tl.constexpr,  # {16, 32, 64, 128, 256, 512}
+        BLOCK_SIZE_D: tl.constexpr,  # {16, 32, 64, 128, 256, 512}
+):
+    pid_b = tl.program_id(0)
+    pid_n = tl.program_id(1)
+    pid_h = tl.program_id(2)
+
+    x_ptr = (x_ptr + pid_b * stride_xb + pid_n * BLOCK_SIZE_N * stride_xn +
+             pid_h * BLOCK_SIZE_H * stride_xh)
+
+    off_n = tl.arange(0, BLOCK_SIZE_N)
+    off_h = tl.arange(0, BLOCK_SIZE_H)
+    off_d = tl.arange(0, BLOCK_SIZE_D)
+
+    cur_block_size_n = min(seq_len - pid_n * BLOCK_SIZE_N, BLOCK_SIZE_N)
+
+    x_mask = ((off_n < seq_len - pid_n * BLOCK_SIZE_N)[:, None, None]
+              & (off_h < num_heads - pid_h * BLOCK_SIZE_H)[None, :, None]
+              & (off_d < head_dim)[None, None, :])
+    x = tl.load(
+        x_ptr + off_n[:, None, None] * stride_xn +
+        off_h[None, :, None] * stride_xh + off_d[None, None, :] * stride_xd,
+        mask=x_mask,
+        other=0,
+    )
+    if pool_type == 0:
+        y = tl.sum(x, axis=0) / cur_block_size_n
+    elif pool_type == 1:
+        y = tl.max(x, axis=0)
+    elif pool_type == 2:
+        y = tl.min(x, axis=0)
+    elif pool_type == 3:
+        y = tl.max(tl.abs(x), axis=0)
+    elif pool_type == 4:
+        y = tl.sum(x, axis=0)
+    else:
+        y = tl.sum(x, axis=0) / cur_block_size_n
+    y_ptr = (y_ptr + pid_b * stride_yb + pid_n * stride_yn +
+             pid_h * BLOCK_SIZE_H * stride_yh)
+    y_mask = (off_h < num_heads -
+              pid_h * BLOCK_SIZE_H)[:, None] & (off_d < head_dim)[None, :]
+    tl.store(y_ptr + off_h[:, None] * stride_yh + off_d[None, :] * stride_yd,
+             y,
+             mask=y_mask)
+
+
+def triton_bnhd_pool(x: torch.Tensor,
+                     kernel_size: int,
+                     pool_type: str = "avg"):
+    b, n, h, d = x.shape
+    assert d in {16, 32, 64, 128}
+    assert kernel_size in {16, 32, 64, 128, 256, 512}
+    m = triton.cdiv(n, kernel_size)
+    y = torch.zeros(b, m, h, d, device=x.device, dtype=x.dtype)
+
+    if pool_type == "last":
+        if n % kernel_size == 0:
+            return x[:, kernel_size - 1::kernel_size, ...]
+        else:
+            return torch.cat(
+                (x[:, kernel_size - 1::kernel_size, ...], x[:, -1:, ...]),
+                dim=1)
+
+    block_size_h = int(triton.next_power_of_2(h))
+    while kernel_size * block_size_h * d > 128 * 128 * 128:
+        block_size_h = int(block_size_h // 2)
+
+    block_size_d = int(triton.next_power_of_2(d))
+    pool_str_to_type = {"avg": 0, "max": 1, "min": 2, "maxabs": 3, "sum": 4}
+    pool_type_index = int(pool_str_to_type[pool_type])
+
+    grid = lambda META: (
+        b,
+        triton.cdiv(n, META["BLOCK_SIZE_N"]),
+        triton.cdiv(h, META["BLOCK_SIZE_H"]),
+    )
+    bnhd_pool_kernel[grid](
+        x,
+        y,
+        pool_type_index,
+        b,
+        n,
+        h,
+        d,
+        x.stride(0),
+        x.stride(1),
+        x.stride(2),
+        x.stride(3),
+        y.stride(0),
+        y.stride(1),
+        y.stride(2),
+        y.stride(3),
+        BLOCK_SIZE_N=kernel_size,
+        BLOCK_SIZE_H=block_size_h,
+        BLOCK_SIZE_D=block_size_d,
+    )
+    return y
+
+
+def torch_bhn_sumpool(x: torch.Tensor, kernel_size: int):
+    b, h, n = x.shape
+    x = torch.nn.functional.pad(
+        x,
+        (
+            0,
+            math.ceil(n / kernel_size) * kernel_size - n,
+        ),
+        value=0,
+    )
+    x = x.view(b, h, -1, kernel_size).sum(-1)
+    return x
+
+
+def score_cover_topk(x: torch.Tensor, score: float):
+    cumsum_x = torch.cumsum(torch.sort(x, dim=-1, descending=True).values,
+                            dim=-1)
+    topk = torch.sum(cumsum_x <= score, dim=-1) + 1
+    return topk
+
+
+def sum_all_diagonal_matrix(mat: torch.tensor):
+    b, h, n, m = mat.shape
+    mat_padded = torch.nn.functional.pad(mat, (n - 1, 0), value=0)
+    mat_strided = mat_padded.as_strided(
+        (b, h, m, n), (h * n * (n + m - 1), n * (n + m - 1), 1, n + m))
+    sum_diags = torch.sum(mat_strided, -1)
+    return sum_diags
+
+
+def get_active_blocks(
+    q,
+    k,
+    v,
+    block_size,
+    gamma,
+    min_budget,
+    max_budget,
+    tau=0,
+    gqa_interleave=False,
+):
+    batch_size, seq_len, num_heads, head_dim = q.shape
+    num_share_q_heads = num_heads // k.shape[2]
+    num_blocks = math.ceil(seq_len / block_size)
+    max_budget = min(max_budget, num_blocks)
+    # last qk attention, qk shape: [batch_size, num_heads, block_size, seq_len]
+    last_q = q[:, -block_size:, :, :] / math.sqrt(head_dim)
+    if not gqa_interleave:
+        qk = torch.einsum(
+            "bihgd, bjhgd -> bhgij",
+            last_q.view(last_q.shape[0], last_q.shape[1], -1,
+                        num_share_q_heads, head_dim),
+            k.view(k.shape[0], k.shape[1], -1, 1, head_dim),
+        )
+    else:
+        qk = torch.einsum(
+            "bihgd, bjhgd -> bhgij",
+            last_q.view(last_q.shape[0], last_q.shape[1], num_share_q_heads,
+                        -1, head_dim),
+            k.view(k.shape[0], k.shape[1], 1, -1, head_dim),
+        )
+    global causal_mask
+    if causal_mask is None:
+        causal_mask = torch.arange(0, block_size, device=last_q.device)
+        causal_mask = causal_mask[:, None] >= causal_mask[None, :]
+        causal_mask = causal_mask[None, None, None, ...]
+    qk[...,
+       -block_size:].masked_fill_(~causal_mask[..., :block_size, :block_size],
+                                  float("-inf"))
+    qk = torch.nn.functional.softmax(qk, dim=-1, dtype=torch.float32)
+    qk = rearrange(qk, "b h g i j -> b (h g) i j")
+    slash = sum_all_diagonal_matrix(qk) / qk.shape[-2]
+    vertical = qk.mean(-2)
+    num_vertical_blocks = score_cover_topk(vertical, gamma) // 128 + 1
+    num_slash_blocks = score_cover_topk(slash, gamma) // 128 + 1
+    num_vertical_blocks[num_vertical_blocks < min_budget] = min_budget
+    num_vertical_blocks[num_vertical_blocks > max_budget] = max_budget
+    num_slash_blocks[num_slash_blocks < min_budget] = min_budget
+    num_slash_blocks[num_slash_blocks > max_budget] = max_budget
+    vertical = torch_bhn_sumpool(vertical, block_size)
+    slash = torch_bhn_sumpool(slash, block_size)
+    if not gqa_interleave:
+        avg_k = triton_bnhd_pool(k, block_size).repeat_interleave(
+            num_share_q_heads, 2)
+    else:
+        avg_k = triton_bnhd_pool(k, block_size).repeat(1, 1, num_share_q_heads,
+                                                       1)
+    avg_qk = torch.einsum("bihd, bjhd -> bhij", last_q.mean(1, keepdim=True),
+                          avg_k).squeeze(2)
+    avg_qk = torch.softmax(avg_qk, dim=-1, dtype=torch.float32)
+    kl_div = square_root_js_divergence(avg_qk, vertical)
+    block_sparse_mask = kl_div < tau
+    num_vertical_blocks[block_sparse_mask] = min_budget
+    num_slash_blocks[block_sparse_mask] = min_budget
+    vertical[..., :1] = torch.inf
+    slash[..., -1:] = torch.inf
+    num_slash_blocks = num_slash_blocks.view(batch_size * num_heads)
+    slash = slash.view(batch_size * num_heads, -1)
+    slash_topk = (num_blocks - 1) - slash.topk(
+        min(num_slash_blocks.max().item(), num_blocks), -1).indices
+    slash_topk[torch.arange(slash_topk.shape[-1],
+                            device=num_slash_blocks.device)[None, :] >=
+               num_slash_blocks[:, None]] = 0
+    slash_topk = slash_topk.view(batch_size, num_heads, -1)
+    # get vertical topk
+    num_vertical_blocks = num_vertical_blocks.view(batch_size * num_heads)
+    vertical = vertical.view(batch_size * num_heads, -1)
+    vertical_topk = vertical.topk(
+        min(num_vertical_blocks.max().item(), num_blocks), -1).indices
+    vertical_topk[torch.arange(vertical_topk.shape[-1],
+                               device=num_vertical_blocks.device)[None, :] >=
+                  num_vertical_blocks[:, None]] = 0
+    vertical_topk = vertical_topk.view(batch_size, num_heads, -1)
+    # transform vertical slash index
+    block_idx = transform_veritcal_slash_idx(vertical_topk, slash_topk,
+                                             num_blocks)
+    # get block sparse topk
+    block_causal_mask = None
+    for b, h in block_sparse_mask.nonzero():
+        if block_causal_mask is None:
+            block_causal_mask = torch.tril(
+                torch.ones(num_blocks,
+                           num_blocks,
+                           device=q.device,
+                           dtype=torch.bool))
+        pad_q = math.ceil(seq_len / block_size) * block_size - seq_len
+        avg_q = (torch.nn.functional.pad(q[b, :, h, :], (0, 0, 0, pad_q),
+                                         value=0).view(num_blocks, block_size,
+                                                       head_dim).mean(1))
+        avg_q[-1, :] = avg_q[-1, :] * block_size / (block_size - pad_q)
+        attn = torch.einsum("id, jd -> ij", avg_q / math.sqrt(head_dim),
+                            avg_k[b, :,
+                                  h, :]).masked_fill_(~block_causal_mask,
+                                                      float("-inf"))
+        attn = torch.softmax(attn, dim=-1, dtype=torch.float32).view(-1)
+        block_topk = score_cover_idx(attn, gamma * num_blocks)
+        block_idx[b][h] = torch.unique(
+            torch.cat((block_idx[b][h], block_topk), dim=-1))
+    return block_idx
+
+
+def sparse_prefill_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    gamma: float = 0.9,
+    tau: float = 0.1,
+    min_budget: int = 512,
+    max_budget: int = 32768,
+    gqa_interleave: bool = False,
+    softmax_scale: Optional[float] = 1.0,
+    block_size: int = 128,
+    return_computational_ratio: bool = False,
+) -> torch.Tensor:
+
+    batch_size, q_len, num_q_heads, head_dim = q.shape
+    batch_size, k_len, num_kv_heads, head_dim = k.shape
+    assert batch_size == 1, "currently only support batch_size == 1"
+
+    if q_len == 1:
+        attn_out = flash_attn_func(q, k, v, softmax_scale=softmax_scale)
+        return attn_out
+    assert q.shape[1] == k.shape[1]
+    assert head_dim in {16, 32, 64, 128}
+    assert block_size in {16, 32, 64, 128}
+    min_budget = 1 if min_budget is None else min_budget
+    max_budget = 2147483647 if max_budget is None else max_budget
+    if q_len <= max(2 * block_size,
+                    math.ceil(min_budget / block_size) * block_size):
+        attn_out = flash_attn_func(q,
+                                   k,
+                                   v,
+                                   softmax_scale=softmax_scale,
+                                   causal=True)
+        return attn_out
+    # get vertical slash index
+    block_idx = get_active_blocks(
+        q,
+        k,
+        v,
+        block_size,
+        gamma,
+        math.ceil(min_budget / block_size),
+        math.ceil(max_budget / block_size),
+        tau,
+        gqa_interleave,
+    )
+    attn_out = triton_block_wise_attention(
+        q,
+        k,
+        v,
+        block_idx,
+        block_size,
+        softmax_scale=softmax_scale,
+        gqa_interleave=gqa_interleave,
+    )
+    return attn_out
diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py
index 5c8781b50d2c..b5b48cb040fb 100644
--- a/vllm/entrypoints/cli/serve.py
+++ b/vllm/entrypoints/cli/serve.py
@@ -4,6 +4,7 @@
 
 import uvloop
 
+import vllm
 from vllm.entrypoints.cli.types import CLISubcommand
 from vllm.entrypoints.openai.api_server import run_server
 from vllm.entrypoints.openai.cli_args import (make_arg_parser,
@@ -24,6 +25,10 @@ def cmd(args: argparse.Namespace) -> None:
         if hasattr(args, 'model_tag') and args.model_tag is not None:
             args.model = args.model_tag
 
+        if vllm.envs.VLLM_USE_SP_PREFILL:
+            args.enable_chunked_prefill = False
+            args.max_num_seqs = 1
+
         uvloop.run(run_server(args))
 
     def validate(self, args: argparse.Namespace) -> None:
diff --git a/vllm/envs.py b/vllm/envs.py
index f80bf878f79c..efed170199bb 100644
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -11,6 +11,7 @@
     VLLM_PORT: Optional[int] = None
     VLLM_RPC_BASE_PATH: str = tempfile.gettempdir()
     VLLM_USE_MODELSCOPE: bool = False
+    VLLM_USE_SP_PREFILL: bool = False
     VLLM_RINGBUFFER_WARNING_INTERVAL: int = 60
     VLLM_NCCL_SO_PATH: Optional[str] = None
     LD_LIBRARY_PATH: Optional[str] = None
@@ -223,6 +224,8 @@ def maybe_convert_int(value: Optional[str]) -> Optional[int]:
     # note that the value is true or false, not numbers
     "VLLM_USE_MODELSCOPE":
     lambda: os.environ.get("VLLM_USE_MODELSCOPE", "False").lower() == "true",
+    "VLLM_USE_SP_PREFILL":
+    lambda: os.environ.get("VLLM_USE_SP_PREFILL", "False").lower() == "true",
 
     # Interval in seconds to log a warning message when the ring buffer is full
     "VLLM_RINGBUFFER_WARNING_INTERVAL":


Reviews:
user: login: rainj-me
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-16T20:13:24Z
user: login: artetaout
state: COMMENTED
author_association: NONE
submitted_at: 2025-04-17T03:02:59Z
user: login: artetaout
state: COMMENTED
author_association: NONE
submitted_at: 2025-04-17T03:03:30Z

Review comments:
diff_hunk: @@ -34,6 +34,16 @@
     from vllm.worker.model_runner import (ModelInputForGPUBuilder,
                                           ModelInputForGPUWithSamplingMetadata)
 
+import vllm.envs
+import vllm.attention.backends.sparse_prefill as sparse_prefill
+
+cfg = {
+    'prefill_gamma': 0.9,
+    'prefill_tau': 0.1,
+    'prefill_min_budget': 512,
+    'prefill_max_budget': 32768,
path: vllm/attention/backends/flash_attn.py
user: login: rainj-me
body: Are these values potentially update in future?
created_at: 2025-04-16T20:08:08Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,1140 @@
+import torch
+import triton
+import triton.language as tl
+from typing import List, Optional, Tuple, Union
+import math
+from einops import rearrange
+causal_mask = None
+
+def gpu_info():
+    if torch.cuda.is_available():
+        device_name = torch.cuda.get_device_name(0).lower()
+        device_capability = torch.cuda.get_device_capability()
+        major, minor = device_capability
+        return device_name, major
+    return None, None
+
+
+GPU_NAME, GPU_MAJOR = gpu_info()
+
+def get_num_warps_stages(head_dim, block_size, gpu_name):
+    gpu_name = gpu_name.lower()
+    # Determine if head_dim and block_size exceed 64
+    head_large = head_dim > 64
+    block_large = block_size > 64
+
+    if "h100" in gpu_name:
path: vllm/attention/backends/sparse_prefill.py
user: login: rainj-me
body: Same here, just double confirm whether we need to customize these values.
created_at: 2025-04-16T20:08:58Z
author_association: COLLABORATOR
diff_hunk: @@ -34,6 +34,16 @@
     from vllm.worker.model_runner import (ModelInputForGPUBuilder,
                                           ModelInputForGPUWithSamplingMetadata)
 
+import vllm.envs
+import vllm.attention.backends.sparse_prefill as sparse_prefill
+
+cfg = {
+    'prefill_gamma': 0.9,
+    'prefill_tau': 0.1,
+    'prefill_min_budget': 512,
+    'prefill_max_budget': 32768,
path: vllm/attention/backends/flash_attn.py
user: login: artetaout
body: yes, plan to make them configurable in the future
created_at: 2025-04-17T03:02:59Z
author_association: NONE
diff_hunk: @@ -0,0 +1,1140 @@
+import torch
+import triton
+import triton.language as tl
+from typing import List, Optional, Tuple, Union
+import math
+from einops import rearrange
+causal_mask = None
+
+def gpu_info():
+    if torch.cuda.is_available():
+        device_name = torch.cuda.get_device_name(0).lower()
+        device_capability = torch.cuda.get_device_capability()
+        major, minor = device_capability
+        return device_name, major
+    return None, None
+
+
+GPU_NAME, GPU_MAJOR = gpu_info()
+
+def get_num_warps_stages(head_dim, block_size, gpu_name):
+    gpu_name = gpu_name.lower()
+    # Determine if head_dim and block_size exceed 64
+    head_large = head_dim > 64
+    block_large = block_size > 64
+
+    if "h100" in gpu_name:
path: vllm/attention/backends/sparse_prefill.py
user: login: artetaout
body: yes, plan to make them configurable
created_at: 2025-04-17T03:03:30Z
author_association: NONE

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-16T06:33:37Z
body: 👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run `fastcheck` CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your `fastcheck` build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping `simon-mo` or `khluu` to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add `ready` label to the PR or enable auto-merge.

🚀
author_association: NONE
user: login: artetaout
created_at: 2025-04-16T06:44:20Z
body: * vanilla
```
============ Serving Benchmark Result ============
Successful requests:                     10        
Benchmark duration (s):                  273.04    
Total input tokens:                      1000000   
Total generated tokens:                  100       
Request throughput (req/s):              0.04      
Output token throughput (tok/s):         0.37      
Total Token throughput (tok/s):          3662.81   
---------------Time to First Token----------------
Mean TTFT (ms):                          49882.45  
Median TTFT (ms):                        51862.56  
P99 TTFT (ms):                           55200.95  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          249.45    
Median TPOT (ms):                        275.36    
P99 TPOT (ms):                           278.80    
---------------Inter-token Latency----------------
Mean ITL (ms):                           249.45    
Median ITL (ms):                         266.17    
P99 ITL (ms):                            336.12    
==================================================
```
* with sparse-prefill
```
============ Serving Benchmark Result ============
Successful requests:                     10        
Benchmark duration (s):                  157.29    
Total input tokens:                      1000000   
Total generated tokens:                  100       
Request throughput (req/s):              0.06      
Output token throughput (tok/s):         0.64      
Total Token throughput (tok/s):          6358.48   
---------------Time to First Token----------------
Mean TTFT (ms):                          29726.09  
Median TTFT (ms):                        30982.87  
P99 TTFT (ms):                           33067.83  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          18.25     
Median TPOT (ms):                        18.20     
P99 TPOT (ms):                           18.66     
---------------Inter-token Latency----------------
Mean ITL (ms):                           18.25     
Median ITL (ms):                         17.67     
P99 ITL (ms):                            27.79     
==================================================
```
author_association: NONE