Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / attention_residuals.py: 0%

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"""Block attention residuals with eager math and the ops-transformer fused API.""" 

16 

17import importlib 

18from collections.abc import Callable 

19 

20# Torch is the native framework; the legacy platform-agnostic lint rule is obsolete here. 

21# pylint: disable=forbidden-backend-import 

22import torch 

23 

24 

25def _block_attention_residuals_eager(prefix: torch.Tensor, history: torch.Tensor, 

26 projection: torch.Tensor, norm_weight: torch.Tensor, 

27 norm_eps: float) -> torch.Tensor: 

28 values = torch.cat((history, prefix.unsqueeze(1)), dim=1) 

29 values_fp32 = values.float() 

30 normalized = values_fp32 * torch.rsqrt(values_fp32.square().mean(-1, keepdim=True) + norm_eps) 

31 query = projection.float().reshape(-1) * norm_weight.float() 

32 probabilities = (normalized * query).sum(-1).softmax(-1) 

33 return (probabilities.unsqueeze(-1) * values_fp32).sum(1).to(values.dtype) 

34 

35 

36def _load_block_attention_residuals() -> Callable: 

37 """Resolve the differentiable ops-transformer API and its backward dependency.""" 

38 try: 

39 # CANN is optional for CPU and eager-only callers. 

40 ops = importlib.import_module("cann_ops_transformer") 

41 except (ImportError, OSError) as error: 

42 raise ImportError( 

43 "AttnRes backend='ascendc' could not load cann_ops_transformer. Install the standalone " 

44 "ops-transformer Python package with compatible PyTorch, torch-npu and CANN libraries. " 

45 f"See hyper_parallel/components/functional/README.md. Original error: {error}" 

46 ) from error 

47 block_attention_residuals_op = getattr(ops, "block_attention_residuals", None) 

48 block_attention_residuals_backward = getattr(ops, "block_attention_residuals_backward", None) 

49 if not callable(block_attention_residuals_op) or not callable(block_attention_residuals_backward): 

50 raise ImportError( 

51 "AttnRes backend='ascendc' requires cann_ops_transformer.block_attention_residuals and " 

52 "block_attention_residuals_backward. Install a package containing both training APIs " 

53 "and build block_attention_residuals plus block_attention_residuals_grad. " 

54 "See hyper_parallel/components/functional/README.md." 

55 ) 

56 return block_attention_residuals_op 

57 

58 

59def _block_attention_residuals_fused(prefix: torch.Tensor, history: torch.Tensor, 

60 projection: torch.Tensor, norm_weight: torch.Tensor, 

61 norm_eps: float) -> torch.Tensor: 

62 """Prepare differentiable inputs and invoke the fused block attention residuals API.""" 

63 block_attention_residuals_op = _load_block_attention_residuals() 

64 # Native forward accepts [H], but its training backward requires [1,H]. 

65 if projection.ndim == 1: 

66 projection = projection.unsqueeze(0) 

67 inputs = (prefix, history, projection, norm_weight) 

68 common_dtype = prefix.dtype 

69 for tensor in inputs[1:]: 

70 common_dtype = torch.promote_types(common_dtype, tensor.dtype) 

71 # The native API requires one dtype; casts must preserve all four gradient edges. 

72 prepared = tuple(tensor.to(dtype=common_dtype).contiguous() for tensor in inputs) 

73 result = block_attention_residuals_op(*prepared, valid_block_num=-1, norm_eps=norm_eps) 

74 return result.to(torch.promote_types(prefix.dtype, history.dtype)) 

75 

76 

77def block_attention_residuals(prefix: torch.Tensor, history: torch.Tensor, 

78 projection: torch.Tensor, norm_weight: torch.Tensor, *, norm_eps: float, 

79 backend: str = "eager") -> torch.Tensor: 

80 """Mix previous blocks and the current residual through depth attention. 

81 

82 Args: 

83 prefix: Current residual of shape [T,H], with complete hidden H. 

84 history: Previous blocks of shape [T,N,H]. 

85 projection: Learned pseudo-query of shape [H] or [1,H]. 

86 norm_weight: Learned RMS normalization scale of shape [H]. 

87 norm_eps: RMSNorm epsilon from the calling model layer. 

88 backend: ``eager`` uses FP32 intermediates. ``ascendc`` calls the 

89 differentiable ops-transformer API; native errors propagate unchanged. 

90 

91 Returns: 

92 Tensor of shape [T,H] with the promoted prefix/history dtype. Mixed 

93 inputs are differentiably promoted to one dtype for the fused operator. 

94 Empty T or N retains zero gradients without a native launch. 

95 

96 Raises: 

97 ImportError: The optional CANN package cannot load or lacks a training API. 

98 

99 Note: 

100 Models own block boundaries and learned parameters. The fused operator 

101 supports first-order gradients. Its dispatcher and native entry validate 

102 device and shape limits; kernel execution errors propagate unchanged. 

103 """ 

104 if backend == "eager": 

105 return _block_attention_residuals_eager(prefix, history, projection, norm_weight, norm_eps) 

106 if backend != "ascendc": 

107 raise ValueError("AttnRes backend must be 'eager' or 'ascendc'") 

108 if prefix.shape[0] == 0 or history.shape[1] == 0: 

109 # Empty views keep zero gradients without inspecting unused values. 

110 zero = sum(tensor.reshape(-1)[:0].sum() for tensor in (history, projection, norm_weight)) 

111 dtype = torch.promote_types(prefix.dtype, history.dtype) 

112 return prefix.to(dtype) + zero.to(dtype) 

113 return _block_attention_residuals_fused(prefix, history, projection, norm_weight, norm_eps)