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

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"""NoPE MLA with a required sigmoid output gate.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Mapping 

20from typing import Any 

21 

22# The high-performance attention modules use native Torch. 

23import torch # pylint: disable=forbidden-backend-import 

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

25 

26from hyper_parallel.components.modules.mla_attention import MLAAttention 

27 

28 

29class GatedMLAAttention(MLAAttention): 

30 """Define NoPE attention and its mandatory output gate in one math module. 

31 

32 The shared key features retain the source projection/checkpoint dimensions; 

33 NoPE disables their rotation. Model adapters only preserve forward APIs. 

34 """ 

35 

36 def __init__(self, *, module: nn.Module, module_fqn: str = "", 

37 context: Mapping[str, Any] | None = None) -> None: 

38 """Validate the source math and retain its managed projection modules.""" 

39 if not getattr(module, "use_nope", False): 

40 raise ValueError("GatedMLAAttention requires a NoPE source module") 

41 if not getattr(module, "use_output_gate", False): 

42 raise ValueError("GatedMLAAttention requires an enabled output gate") 

43 if type(module.g_proj) is not nn.Linear: # pylint: disable=unidiomatic-typecheck 

44 raise ValueError("Gated MLA requires a plain Linear output gate") 

45 super().__init__(module=module, module_fqn=module_fqn, context=context) 

46 self.use_nope = True 

47 self.g_proj = module.g_proj 

48 

49 @staticmethod 

50 def apply_nope(tensor: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None, 

51 *, interleaved: bool = False) -> torch.Tensor: 

52 """Keep Q/shared-K features unchanged through the CP rotary callback. 

53 

54 Args: 

55 tensor: Query or shared-key features in the runtime's BSND layout. 

56 position_embeddings: Ignored by NoPE attention. 

57 interleaved: Ignored because no positional rotation is applied. 

58 

59 Returns: 

60 The original tensor, without a copy or a positional transform. 

61 """ 

62 del position_embeddings, interleaved 

63 return tensor 

64 

65 def _project_attention_inputs( 

66 self, hidden_states: torch.Tensor, 

67 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None, past_key_values: Any | None, 

68 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 

69 """Use the common projections while enforcing the NoPE math locally.""" 

70 del position_embeddings 

71 return super()._project_attention_inputs(hidden_states, None, past_key_values) 

72 

73 def project_output(self, attention_output: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: 

74 """Apply the source gate before Wo, without an additional output Norm. 

75 

76 Args: 

77 attention_output: Flattened token-local attention output. 

78 hidden_states: Original token-local input for the gate projection. 

79 

80 Returns: 

81 Gated and projected hidden states. 

82 """ 

83 attention_output = attention_output * self.g_proj(hidden_states).sigmoid() 

84 return self.o_proj(attention_output)