Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / __init__.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"""Public high-performance function interfaces. 

16 

17Exports are resolved lazily (PEP 562): most functions wrap NPU-only kernels 

18whose modules import ``torch_npu`` at top level, so importing this package 

19must not force every backend onto CPU-only consumers that only need a 

20single submodule (e.g. ``functional.npu_grouped_swiglu``). 

21""" 

22 

23# PEP 562 requires the special ``__getattr__`` name, and the public symbols 

24# listed in ``__all__`` are materialized dynamically by that hook. 

25# pylint: disable=invalid-name,undefined-all-variable 

26 

27import importlib 

28from typing import Any 

29 

30_EXPORT_TO_MODULE = { 

31 "aggregate_hidden": "aggregate_hidden", 

32 "apply_rotary_pos_emb": "rotary_embedding", 

33 "apply_rotary_pos_emb_interleave": "rotary_embedding", 

34 "apply_rotary_pos_emb_single": "rotary_embedding", 

35 "apply_rotary_pos_emb_single_freqs": "rotary_embedding", 

36 "attention_rescale": "attention_rescale", 

37 "aux_loss_auto_scale": "aux_loss", 

38 "aux_loss_scale_context": "aux_loss", 

39 "bind_aux_loss_scale": "aux_loss", 

40 "dsa_indexer": "dsa_indexer", 

41 "dsa_kl_loss": "dsa_kl_loss", 

42 "dsa_sparse_attention": "dsa_sparse_attention", 

43 "dsa_sparse_attention_rescale": "dsa_sparse_attention_rescale", 

44 "grouped_matmul": "grouped_matmul", 

45 "chunk_gated_delta_rule": "gated_delta_net", 

46 "fused_chunk_kda": "kimi_delta_attention", 

47 "fused_chunk_kda_p2p": "kimi_delta_attention", 

48 "mhc_post": "mhc_post", 

49 "mhc_pre": "mhc_pre", 

50 "moe_token_permute": "moe_token_permute", 

51 "moe_token_unpermute": "moe_token_unpermute", 

52 "npu_fusion_attention_forward": "npu_fusion_attention", 

53 "npu_grouped_swiglu": "npu_grouped_swiglu", 

54 "rms_norm": "rms_norm", 

55 "set_aux_loss_scale": "aux_loss", 

56 "sink_attention": "sink_attention", 

57 "sinkhorn": "sinkhorn", 

58 "swiglu": "swiglu", 

59} 

60 

61 

62def __getattr__(name: str) -> Any: 

63 """Resolve a public function by importing its owning submodule lazily.""" 

64 submodule = _EXPORT_TO_MODULE.get(name) 

65 if submodule is None: 

66 raise AttributeError( 

67 f"module {__name__!r} has no attribute {name!r}" 

68 ) 

69 module = importlib.import_module( 

70 f"hyper_parallel.components.functional.{submodule}" 

71 ) 

72 value = getattr(module, name) 

73 globals()[name] = value 

74 return value 

75 

76 

77__all__ = [ 

78 "aggregate_hidden", 

79 "apply_rotary_pos_emb", 

80 "apply_rotary_pos_emb_interleave", 

81 "apply_rotary_pos_emb_single", 

82 "apply_rotary_pos_emb_single_freqs", 

83 "attention_rescale", 

84 "aux_loss_auto_scale", 

85 "aux_loss_scale_context", 

86 "bind_aux_loss_scale", 

87 "dsa_indexer", 

88 "dsa_kl_loss", 

89 "dsa_sparse_attention", 

90 "dsa_sparse_attention_rescale", 

91 "grouped_matmul", 

92 "chunk_gated_delta_rule", 

93 "fused_chunk_kda", 

94 "fused_chunk_kda_p2p", 

95 "mhc_post", 

96 "mhc_pre", 

97 "moe_token_permute", 

98 "moe_token_unpermute", 

99 "npu_fusion_attention_forward", 

100 "npu_grouped_swiglu", 

101 "rms_norm", 

102 "set_aux_loss_scale", 

103 "sink_attention", 

104 "sinkhorn", 

105 "swiglu", 

106]