chore: import upstream snapshot with attribution

2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
@@ -0,0 +1,97 @@
+# 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