131 lines
5.6 KiB
Python
131 lines
5.6 KiB
Python
################################################################################
|
|
# Copyright(c)2020-2025 Shanghai Biren Technology Co., Ltd. All rights reserved.
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
################################################################################
|
|
|
|
from typing import Optional
|
|
|
|
import torch
|
|
|
|
import vllm.attention.layer
|
|
from vllm.attention.layer import (maybe_save_kv_layer_to_connector,
|
|
wait_for_kv_layer_from_connector)
|
|
from vllm.forward_context import ForwardContext, get_forward_context
|
|
|
|
#direct_register_custom_op(
|
|
# op_name="unified_attention",
|
|
# op_func=unified_attention,
|
|
# mutates_args=[],
|
|
# fake_impl=unified_attention_fake,
|
|
# dispatch_key=current_platform.dispatch_key,
|
|
#)
|
|
|
|
#direct_register_custom_op(
|
|
# op_name="unified_attention_with_output",
|
|
# op_func=unified_attention_with_output,
|
|
# mutates_args=["output"],
|
|
# fake_impl=unified_attention_with_output_fake,
|
|
# dispatch_key=current_platform.dispatch_key,
|
|
#)
|
|
|
|
|
|
def forward_(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
# For some alternate attention backends like MLA the attention output
|
|
# shape does not match the query shape, so we optionally let the model
|
|
# definition specify the output tensor shape.
|
|
output_shape: Optional[torch.Size] = None,
|
|
) -> torch.Tensor:
|
|
"""
|
|
The KV cache is stored inside this class and is accessed via
|
|
`self.kv_cache`.
|
|
|
|
Attention metadata (`attn_metadata`) is set using a context manager in
|
|
the model runner's `execute_model` method. It is accessed via forward
|
|
context using
|
|
`vllm.forward_context.get_forward_context().attn_metadata`.
|
|
"""
|
|
if self.calculate_kv_scales:
|
|
attn_metadata = get_forward_context().attn_metadata
|
|
if attn_metadata.enable_kv_scales_calculation:
|
|
self.calc_kv_scales(query, key, value)
|
|
if self.use_output:
|
|
output_shape = (output_shape
|
|
if output_shape is not None else query.shape)
|
|
output = torch.empty(output_shape,
|
|
dtype=query.dtype,
|
|
device=query.device)
|
|
hidden_size = output_shape[-1]
|
|
# We skip reshaping query, key and value tensors for the MLA
|
|
# backend since these tensors have different semantics and are
|
|
# processed differently.
|
|
if not self.use_mla:
|
|
# Reshape the query, key, and value tensors.
|
|
# NOTE(woosuk): We do this outside the custom op to minimize the
|
|
# CPU overheads from the non-CUDA-graph regions.
|
|
query = query.view(-1, self.num_heads, self.head_size)
|
|
output = output.view(-1, self.num_heads, self.head_size)
|
|
if key is not None:
|
|
key = key.view(-1, self.num_kv_heads, self.head_size)
|
|
if value is not None:
|
|
value = value.view(-1, self.num_kv_heads, self.head_size)
|
|
if self.use_direct_call:
|
|
forward_context: ForwardContext = get_forward_context()
|
|
attn_metadata = forward_context.attn_metadata
|
|
if isinstance(attn_metadata, dict):
|
|
attn_metadata = attn_metadata[self.layer_name]
|
|
self_kv_cache = self.kv_cache[forward_context.virtual_engine]
|
|
self.impl.forward(self,
|
|
query,
|
|
key,
|
|
value,
|
|
self_kv_cache,
|
|
attn_metadata,
|
|
output=output)
|
|
else:
|
|
torch.ops.vllm.unified_attention_with_output(
|
|
query, key, value, output, self.layer_name)
|
|
return output.view(-1, hidden_size)
|
|
else:
|
|
if self.use_direct_call:
|
|
wait_for_kv_layer_from_connector(self.layer_name)
|
|
forward_context = get_forward_context()
|
|
attn_metadata = forward_context.attn_metadata
|
|
if isinstance(attn_metadata, dict):
|
|
attn_metadata = attn_metadata[self.layer_name]
|
|
self_kv_cache = self.kv_cache[forward_context.virtual_engine]
|
|
output = self.impl.forward(self, query, key, value, self_kv_cache,
|
|
attn_metadata)
|
|
maybe_save_kv_layer_to_connector(self.layer_name, self_kv_cache)
|
|
return output
|
|
else:
|
|
# return torch.ops.vllm.unified_attention(
|
|
# query, key, value, self.layer_name)
|
|
wait_for_kv_layer_from_connector(self.layer_name)
|
|
forward_context = get_forward_context()
|
|
attn_metadata = forward_context.attn_metadata
|
|
if isinstance(attn_metadata, dict):
|
|
attn_metadata = attn_metadata[self.layer_name]
|
|
self_kv_cache = self.kv_cache[forward_context.virtual_engine]
|
|
output = self.impl.forward(self, query, key, value, self_kv_cache,
|
|
attn_metadata)
|
|
maybe_save_kv_layer_to_connector(self.layer_name, self_kv_cache)
|
|
return output
|
|
|
|
|
|
vllm.attention.layer.Attention.forward = forward_
|