Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / gdn_triton.py: 100%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15 

16"""Structure-preserving replacement for the Triton-Ascend GDN primitive.""" 

17 

18from __future__ import annotations 

19 

20import copy 

21from collections.abc import Mapping 

22from importlib import import_module 

23from typing import Any 

24 

25from torch import nn # pylint: disable=forbidden-backend-import 

26 

27from hyper_parallel.models.replacement import module_replacement 

28 

29 

30@module_replacement 

31class TritonGDN: 

32 """Factory selecting the Triton-Ascend GDN primitive for an existing module. 

33 

34 The returned module retains its source type so the model's forward, CP 

35 wrapper, and checkpoint layout remain unchanged. 

36 """ 

37 

38 def __new__( 

39 cls, 

40 *, 

41 module: nn.Module, 

42 module_fqn: str = "", 

43 context: Mapping[str, Any] | None = None, 

44 ) -> nn.Module: 

45 """Bind the full Triton-Ascend GDN forward/backward without changing weights. 

46 

47 Args: 

48 module: GDN module exposing an instance-level chunk primitive. 

49 module_fqn: Source module's fully qualified name. 

50 context: Model replacement context. 

51 

52 Returns: 

53 A shallow copy with its chunk primitive bound to the Triton implementation. 

54 

55 Raises: 

56 TypeError: The source does not expose the expected primitive. 

57 """ 

58 del cls, context 

59 if not callable(getattr(module, "chunk_gated_delta_rule", None)): 

60 raise TypeError(f"{module_fqn}: GDN replacement requires callable chunk_gated_delta_rule") 

61 # Triton-Ascend is optional and should load only when this replacement is selected. 

62 triton_rule = import_module( 

63 "hyper_parallel.components.functional.gated_delta_net" 

64 ).chunk_gated_delta_rule 

65 replacement = copy.copy(module) 

66 replacement.chunk_gated_delta_rule = triton_rule 

67 return replacement