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

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"""Public high-performance module interfaces. 

16 

17Exports are resolved lazily so modules backed only by torch-npu do not import 

18unrelated optional Omni custom operators. Accessing an Omni-backed module still 

19loads and validates that dependency normally. 

20""" 

21 

22__all__: list[str] = [] 

23 

24import importlib 

25from typing import TYPE_CHECKING, Any 

26 

27if TYPE_CHECKING: 

28 from hyper_parallel.components.modules.mtp import ( 

29 MTPLayer, MultiTokenPrediction, MTPOutput, build_mtp, 

30 ) 

31 from hyper_parallel.components.modules.dsa_attention import DSAAttention, DeepseekV32DSAAttention 

32 from hyper_parallel.components.modules.gqa_attention import GatedGQAAttention, GQAAttention 

33 from hyper_parallel.components.modules.grouped_experts import GroupedExperts 

34 from hyper_parallel.components.modules.linear import Linear 

35 from hyper_parallel.components.modules.mhc import MhcPostModule, MhcPreModule 

36 from hyper_parallel.components.modules.mla_attention import MLAAttention 

37 from hyper_parallel.components.modules.rms_norm import OffsetRMSNorm, RMSNorm 

38 from hyper_parallel.components.modules.shared_expert import SharedExpert 

39 from hyper_parallel.components.modules.swiglu_mlp import SwiGLUMLP 

40 

41_EXPORT_TO_MODULE = { 

42 "TritonGDN": "gdn_triton", 

43 "build_mtp": "mtp", 

44 "MTPLayer": "mtp", 

45 "MultiTokenPrediction": "mtp", 

46 "MTPOutput": "mtp", 

47 "DeepseekV32DSAAttention": "dsa_attention", 

48 "DSAAttention": "dsa_attention", 

49 "GQAAttention": "gqa_attention", 

50 "GatedGQAAttention": "gqa_attention", 

51 "EngramModule": "engram", 

52 "GroupedExperts": "grouped_experts", 

53 "Linear": "linear", 

54 "KimiDeltaAttention": "kimi_delta_attention", 

55 "KimiRMSNormGated": "kimi_delta_attention", 

56 "MhcPostModule": "mhc", 

57 "MhcPreModule": "mhc", 

58 "NgramHashMapping": "engram", 

59 "PipelinedMhcModule": "mhc", 

60 "MLAAttention": "mla_attention", 

61 "OffsetRMSNorm": "rms_norm", 

62 "RMSNorm": "rms_norm", 

63 "SharedExpert": "shared_expert", 

64 "SharedCompressedAttentionCPContext": "shared_compressed_dsa_attention", 

65 "SharedCompressedAttentionState": "shared_compressed_dsa_attention", 

66 "SharedCompressedPackedSequence": "shared_compressed_dsa_attention", 

67 "SharedCompressedAttentionTPContext": "shared_compressed_dsa_attention", 

68 "SharedCompressedDSAAttention": "shared_compressed_dsa_attention", 

69 "SharedCompressedDSAIndexer": "shared_compressed_dsa_attention", 

70 "SwiGLUMLP": "swiglu_mlp", 

71} 

72_EXPORT_TO_MODULE.update({name: "module_family" for name in ( 

73 "ModuleFamily", "SelectorSpec", "VariantRule", "FamilyRequest", "ResolvedFamily", 

74 "build_family", "register_module_family", "get_family_build_info", "collect_family_modules", 

75 "FamilyShardingContext", 

76)}) 

77__all__.extend(_EXPORT_TO_MODULE) 

78 

79 

80def __getattr__(name: str) -> Any: # pylint: disable=invalid-name 

81 """Resolve a public class by importing only its owning submodule.""" 

82 submodule = _EXPORT_TO_MODULE.get(name) 

83 if submodule is None: 

84 raise AttributeError(f"module {__name__!r} has no attribute {name!r}") 

85 module = importlib.import_module(f"hyper_parallel.components.modules.{submodule}") 

86 value = getattr(module, name) 

87 globals()[name] = value 

88 return value 

89 

90 

91def __dir__() -> list[str]: 

92 """Include lazy public exports in ``dir()``.""" 

93 return sorted(set(globals()) | set(__all__))