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 "block_attention_residuals": "attention_residuals", 

41 "dsa_indexer": "dsa_indexer", 

42 "dsa_dense_indexer_kl_loss": "dsa_dense", 

43 "dsa_dense_indexer_lse": "dsa_dense", 

44 "dsa_kl_loss": "dsa_kl_loss", 

45 "dsa_sparse_attention": "dsa_sparse_attention", 

46 "dsa_sparse_attention_rescale": "dsa_sparse_attention_rescale", 

47 "grouped_matmul": "grouped_matmul", 

48 "chunk_gated_delta_rule": "gated_delta_net", 

49 "fused_chunk_kda": "kimi_delta_attention", 

50 "fused_chunk_kda_p2p": "kimi_delta_attention", 

51 "mhc_post": "mhc_post", 

52 "mhc_pre": "mhc_pre", 

53 "moe_token_permute": "moe_token_permute", 

54 "moe_token_unpermute": "moe_token_unpermute", 

55 "npu_fusion_attention_forward": "npu_fusion_attention", 

56 "npu_grouped_swiglu": "npu_grouped_swiglu", 

57 "rms_norm": "rms_norm", 

58 "set_aux_loss_scale": "aux_loss", 

59 "sink_attention": "sink_attention", 

60 "sinkhorn": "sinkhorn", 

61 "situ": "activation_situ", 

62 "situ_packed": "activation_situ", 

63 "swiglu": "swiglu", 

64} 

65 

66 

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

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

69 submodule = _EXPORT_TO_MODULE.get(name) 

70 if submodule is None: 

71 raise AttributeError( 

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

73 ) 

74 module = importlib.import_module( 

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

76 ) 

77 value = getattr(module, name) 

78 globals()[name] = value 

79 return value 

80 

81 

82__all__ = [ 

83 "aggregate_hidden", 

84 "apply_rotary_pos_emb", 

85 "apply_rotary_pos_emb_interleave", 

86 "apply_rotary_pos_emb_single", 

87 "apply_rotary_pos_emb_single_freqs", 

88 "attention_rescale", 

89 "aux_loss_auto_scale", 

90 "aux_loss_scale_context", 

91 "bind_aux_loss_scale", 

92 "block_attention_residuals", 

93 "dsa_indexer", 

94 "dsa_dense_indexer_kl_loss", 

95 "dsa_dense_indexer_lse", 

96 "dsa_kl_loss", 

97 "dsa_sparse_attention", 

98 "dsa_sparse_attention_rescale", 

99 "grouped_matmul", 

100 "chunk_gated_delta_rule", 

101 "fused_chunk_kda", 

102 "fused_chunk_kda_p2p", 

103 "mhc_post", 

104 "mhc_pre", 

105 "moe_token_permute", 

106 "moe_token_unpermute", 

107 "npu_fusion_attention_forward", 

108 "npu_grouped_swiglu", 

109 "rms_norm", 

110 "set_aux_loss_scale", 

111 "sink_attention", 

112 "sinkhorn", 

113 "situ", 

114 "situ_packed", 

115 "swiglu", 

116]