Skip to content

GitLab

Commit 5289f7c4 authored by zhuwenwen's avatar zhuwenwen
Browse files

update _custom_ops.py

parent a7e3c9d6
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
vllm/_custom_ops.py
View file @ 5289f7c4
...@@ -255,7 +255,8 @@ def paged_attention_v2_opt_tc( ...@@ -255,7 +255,8 @@ def paged_attention_v2_opt_tc(
max_seq_len: int, max_seq_len: int,
alibi_slopes: Optional[torch.Tensor], alibi_slopes: Optional[torch.Tensor],
kv_cache_dtype: str, kv_cache_dtype: str,
kv_scale: float, k_scale: float,
v_scale: float,
tp_rank: int = 0, tp_rank: int = 0,
blocksparse_local_blocks: int = 0, blocksparse_local_blocks: int = 0,
blocksparse_vert_stride: int = 0, blocksparse_vert_stride: int = 0,
...@@ -265,7 +266,7 @@ def paged_attention_v2_opt_tc( ...@@ -265,7 +266,7 @@ def paged_attention_v2_opt_tc(
torch.ops._C.paged_attention_v2_opt_tc( torch.ops._C.paged_attention_v2_opt_tc(
out, exp_sum, max_logits, tmp_out, query, key_cache, value_cache, 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, num_kv_heads, scale, block_tables, seq_lens, block_size, max_seq_len,
alibi_slopes, kv_cache_dtype, kv_scale, tp_rank, alibi_slopes, kv_cache_dtype, k_scale, v_scale, tp_rank,
blocksparse_local_blocks, blocksparse_vert_stride, blocksparse_local_blocks, blocksparse_vert_stride,
blocksparse_block_size, blocksparse_head_sliding_step) blocksparse_block_size, blocksparse_head_sliding_step)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment