wehub-resource-sync
98 lines
3.5 KiB
Python
98 lines
3.5 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from typing import Any, Dict, Optional, Tuple, Type
|
|
|
|
import torch
|
|
|
|
from ...ragged import RaggedBatchWrapper
|
|
from deepspeed.runtime.config_utils import DeepSpeedConfigModel
|
|
from ..ds_module import DSModuleBase
|
|
from ..module_registry import DSModuleRegistryBase
|
|
from ..configs import DSSelfAttentionConfig
|
|
|
|
|
|
class DSSelfAttentionBase(DSModuleBase):
|
|
"""
|
|
Base mixin for all attention modules. The interface represented by this module
|
|
is broadly:
|
|
|
|
output = attention(query_key_value,
|
|
Optional[kv_cache],
|
|
Optional[attention_mask],
|
|
Optional[attention_bias])
|
|
"""
|
|
|
|
@staticmethod
|
|
def config_class() -> Type[DeepSpeedConfigModel]:
|
|
return DSSelfAttentionConfig
|
|
|
|
def __init__(self, config: DSSelfAttentionConfig, implementation_config: Dict[str, Any]) -> None:
|
|
super().__init__(config, implementation_config)
|
|
|
|
@property
|
|
def kv_block_size(self) -> int:
|
|
"""
|
|
Return preferred granulatity for blocked KV-cache implementation.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
@property
|
|
def q_block_size(self) -> int:
|
|
"""
|
|
Property to calculate blocking granularity for the query dimension.
|
|
This has no impact on the KV-cache structure, but will affect the
|
|
number of attention atoms associated with a batch.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
def build_atoms(self, ragged_batch: RaggedBatchWrapper) -> None:
|
|
"""
|
|
Build the atoms for this module. This is not a strict requirement for the class,
|
|
so this method is a no-op by default rather than abstract.
|
|
"""
|
|
pass
|
|
|
|
def forward(self,
|
|
q_k_v: torch.Tensor,
|
|
kv_cache: torch.Tensor,
|
|
batch: RaggedBatchWrapper,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
attention_bias: Optional[torch.Tensor] = None,
|
|
inv_freqs: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Parameters:
|
|
q_k_v (torch.Tensor): Query, key, and value tensors. Expected shape is:
|
|
[
|
|
batch,
|
|
seq_len,
|
|
2 * self._config.n_heads_kv + self._config.n_heads_q,
|
|
self._config.head_size
|
|
].
|
|
kv_cache (Optional[torch.Tensor]): Key and value cache tensor. Expected shape is
|
|
[
|
|
2,
|
|
batch,
|
|
kv_cache_len,
|
|
self._config.n_heads_kv,
|
|
self._config.head_size
|
|
]. If None, cache is disabled. The `kv_cache_len` dimension does not need to
|
|
be contiguous (it should expand stride by `max_out_tokens`).
|
|
batch (RaggedBatchWrapper): Ragged batch metadata.
|
|
attention_mask (Optional[torch.Tensor]): Attention mask tensor. If None, masking is
|
|
disabled. This will defer to the config in the case of conflicting information.
|
|
This means if the config class is implying causal attention, the mask will be ignored.
|
|
attention_bias (Optional[torch.Tensor]): Attention bias tensor. If None, bias is disabled.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
|
|
class DSSelfAttentionRegistry(DSModuleRegistryBase):
|
|
registry: Dict = {}
|
|
|
|
@staticmethod
|
|
def associated_class() -> Type[DSModuleBase]:
|
|
return DSSelfAttentionBase
|