"vllm/vscode:/vscode.git/clone" did not exist on "2aedb72b217f8961823bbcac488c03b31f8273db"
Commit 668ec4ef authored by zhuwenwen's avatar zhuwenwen
Browse files

Update refactoring operation

parent 8a325c18
......@@ -193,28 +193,28 @@ class PagedAttention:
if envs.VLLM_USE_OPT_OP:
ops.paged_attention_v2_opt(
output,
exp_sums,
max_logits,
tmp_output,
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,
)
output,
exp_sums,
max_logits,
tmp_output,
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,
)
else:
ops.paged_attention_v2(
output,
......
......@@ -36,7 +36,7 @@ class SiluAndMul(CustomOp):
output_shape = (x.shape[:-1] + (d, ))
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
if envs.VLLM_USE_OPT_OP:
ops.silu_and_mul(out, x)
ops.silu_and_mul_opt(out, x)
else:
ops.silu_and_mul(out, x)
return out
......
......@@ -70,11 +70,11 @@ class RMSNorm(CustomOp):
out = torch.empty_like(x)
if envs.VLLM_USE_OPT_OP:
ops.rms_norm_opt(
out,
x,
self.weight.data,
self.variance_epsilon,
)
out,
x,
self.weight.data,
self.variance_epsilon,
)
else:
ops.rms_norm(
out,
......
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