mindspore 2.2.10__cp39-cp39-win_amd64.whl → 2.2.14__cp39-cp39-win_amd64.whl

mindspore/ops/_op_impl/_custom_op/flash_attention/tiling_strategy/strategy.py

@@ -1,67 +0,0 @@
-# Copyright 2023 Huawei Technologies Co., Ltd
-#
-# 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.
-# ============================================================================
-"""The base class of tiling strategy"""
-from abc import ABCMeta
-from abc import abstractmethod
-from collections import namedtuple
-
-TilingPara = namedtuple("TilingPara", "Br last_Br Bc last_Bc Tr Tc")
-
-
-class TilingStrategy(metaclass=ABCMeta):
-    """Tiling strategy interface. All implementations should be defined in this module,
-    otherwise, the UT will fail.
-    """
-
-    _strategies = {}
-
-    def __init__(self, Nq, N, head_dim) -> None:
-        super().__init__()
-        self.Nq = Nq
-        self.N = N
-        self.Br = None
-        self.last_Br = None
-        self.Bc = None
-        self.last_Bc = None
-        self.Tr = None
-        self.Tc = None
-        self.d = head_dim
-
-    def __init_subclass__(cls, **kwargs):
-        TilingStrategy._strategies[cls.strategy_name()] = cls
-
-    @classmethod
-    @abstractmethod
-    def strategy_name(cls):
-        """strategy name"""
-        raise NotImplementedError
-
-    @classmethod
-    def from_strategy_name(cls, stgy_name: str):
-        """from strategy name"""
-        stgy_clz = TilingStrategy._strategies.get(stgy_name)
-        if stgy_clz is None:
-            raise Exception(f"Strategy:{stgy_name} not supported")
-
-        return stgy_clz
-
-    @abstractmethod
-    def tiling(self) -> TilingPara:
-        """tiling"""
-        raise NotImplementedError
-
-    def gen_tiling_para(self) -> TilingPara:
-        """gen tiling para"""
-        return TilingPara(self.Br, self.last_Br, self.Bc, self.last_Bc, self.Tr, self.Tc)

mindspore/ops/_op_impl/_custom_op/flash_attention/tiling_strategy/wukong_tiling.py

@@ -1,62 +0,0 @@
-# Copyright 2023 Huawei Technologies Co., Ltd
-#
-# 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.
-# ============================================================================
-"""wukong tiling"""
-from mindspore.ops._op_impl._custom_op.flash_attention.tiling_strategy.strategy import TilingPara
-from mindspore.ops._op_impl._custom_op.flash_attention.tiling_strategy.strategy import TilingStrategy
-
-
-class WukongTiling(TilingStrategy):
-    """A tiling strategy implementation for wukonghuahua model shape"""
-
-    @classmethod
-    def strategy_name(cls):
-        return "wukong"
-
-    def tiling(self) -> TilingPara:
-        """
-        反向的空间分布待详细分析
-        N = (4096, 1024, 256, 64) 或 77
-        Nq = (4096, 1024, 256, 64)
-        d = dv = (40, 80, 160,  160)
-        """
-        if self.N <= 77:  # [77, 64]
-            # cross-attention or self-attention of (64, 64, 160)
-            self.Bc = self.N
-            self.Tc = self.N // self.Bc
-            if self.d <= 80:  # [40, 80]
-                # 内存瓶颈为在ub中对P*V结果[Br, dv]进行cast
-                self.Br = min(self.Nq, 64)
-                self.Tr = self.Nq // self.Br
-            else:
-                self.Br = min(self.Nq, 64)
-                self.Tr = self.Nq // self.Br
-        else:
-            # self-attention
-            if self.N == 256:
-                self.Bc = 64
-                self.Tc = 1
-                # 内存瓶颈为在ub中对Q*K的结果[Br, Bc]进行cast
-                self.Br = 64
-                self.Tr = self.Nq // self.Br
-            else:
-                self.Bc = 64
-                self.Tc = self.N // self.Bc
-                self.Br = 64
-                self.Tr = self.Nq // self.Br
-
-        self.last_Br = self.Br
-        self.last_Bc = self.Bc
-
-        return self.gen_tiling_para()