Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / gated_delta_net_ascendc.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"""Optional Ascend C GDN core with effective scalar gate and beta inputs.""" 

16 

17from __future__ import annotations 

18 

19import importlib 

20import inspect 

21from collections.abc import Sequence 

22from functools import lru_cache 

23from typing import Any 

24 

25import torch 

26from torch.autograd.function import once_differentiable 

27 

28 

29@lru_cache(maxsize=1) 

30def _load_ops(): 

31 """Keep the optional native runtime out of ordinary framework imports.""" 

32 try: 

33 ops = importlib.import_module("fla_npu.ops.ascendc") 

34 except ImportError as error: 

35 raise RuntimeError("GDN backend='ascendc' requires a compatible fla-npu installation.") from error 

36 required = { 

37 "chunk_gated_delta_rule_fwd": ( 

38 "layout", "use_exp2", "cu_seqlens", "chunk_indices", "disable_recompute", 

39 "use_qk_l2norm_in_kernel", "use_beta_sigmoid_in_kernel", "use_gate_in_kernel", 

40 ), 

41 "chunk_gated_delta_rule_bwd": ( 

42 "layout", "use_exp2", "cu_seqlens", "chunk_indices", "use_qk_l2norm_in_kernel", 

43 "use_beta_sigmoid_in_kernel", "use_gate_in_kernel", 

44 ), 

45 } 

46 for name, parameters in required.items(): 

47 try: 

48 signature = inspect.signature(getattr(ops, name, None)) 

49 except (TypeError, ValueError) as error: 

50 raise RuntimeError(f"fla-npu lacks a compatible public {name} API.") from error 

51 missing = set(parameters).difference(signature.parameters) 

52 if missing: 

53 raise RuntimeError(f"fla-npu {name} lacks required parameters {sorted(missing)}.") 

54 return ops 

55 

56 

57def _packed_offsets(cu_seqlens, tokens): 

58 """Freeze validated document offsets without a device-to-host transfer.""" 

59 if isinstance(cu_seqlens, torch.Tensor): 

60 if (cu_seqlens.device.type != "cpu" or cu_seqlens.ndim != 1 

61 or cu_seqlens.dtype not in (torch.int32, torch.int64)): 

62 raise ValueError("Ascend C GDN cu_seqlens must be a one-dimensional CPU integer tensor or sequence.") 

63 cu_seqlens = cu_seqlens.tolist() 

64 if not isinstance(cu_seqlens, Sequence): 

65 raise ValueError("Ascend C GDN cu_seqlens must be a CPU integer tensor or sequence.") 

66 offsets = tuple(cu_seqlens) 

67 if (len(offsets) < 2 or offsets[0] != 0 or offsets[-1] != tokens 

68 or any(not isinstance(value, int) or isinstance(value, bool) for value in offsets) 

69 or any(start >= end for start, end in zip(offsets, offsets[1:]))): 

70 raise ValueError("Ascend C GDN cu_seqlens must strictly increase from zero to the packed token count.") 

71 return offsets 

72 

73 

74def _packed_metadata(cu_seqlens, tokens): 

75 """Derive canonical native chunk indices from immutable host offsets.""" 

76 if cu_seqlens is None: 

77 return {} 

78 offsets = _packed_offsets(cu_seqlens, tokens) 

79 chunks = tuple(component 

80 for sequence, (start, end) in enumerate(zip(offsets, offsets[1:])) 

81 for chunk in range((end - start + 63) // 64) 

82 for component in (sequence, chunk)) 

83 return {"cu_seqlens": offsets, "chunk_indices": chunks} 

84 

85 

86class _AscendCGDN(torch.autograd.Function): 

87 """Own each invocation's native backward intermediates independently.""" 

88 

89 @staticmethod 

90 def forward(ctx: Any, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, 

91 gate: torch.Tensor, beta: torch.Tensor, metadata: dict) -> torch.Tensor: 

92 """Save only the tensors consumed by the native composite backward.""" 

93 scale = query.shape[-1] ** -0.5 

94 layout = "TND" if metadata else "BSND" 

95 result = _load_ops().chunk_gated_delta_rule_fwd( 

96 query, key, value, gate, beta, scale=scale, layout=layout, chunk_size=64, 

97 initial_state=None, output_final_state=False, disable_recompute=True, 

98 return_intermediate_states=False, state_v_first=False, use_exp2=False, 

99 use_qk_l2norm_in_kernel=False, use_beta_sigmoid_in_kernel=False, 

100 use_gate_in_kernel=False, **metadata, 

101 ) 

102 # The public ABI keeps ten slots even when optional norm/state outputs 

103 # are disabled: o, final_state, g_cumsum, A, beta_eff, h, q_hat, 

104 # k_hat, q_rstd, k_rstd. This wrapper only consumes slots 0, 2 and 3. 

105 if len(result) != 10 or any(not isinstance(result[index], torch.Tensor) for index in (0, 2, 3)): 

106 raise RuntimeError("fla-npu GDN forward must return ten slots with output, g_cumsum and A tensors.") 

107 output, _, gate_cumsum, matrix_a, *_ = result 

108 # Gate and beta transformations remain in the caller's autograd graph. 

109 # Native backward consumes cumulative gate but returns the raw-gate gradient. 

110 ctx.save_for_backward(query, key, value, gate_cumsum, beta, matrix_a) 

111 ctx.native_metadata, ctx.layout, ctx.scale = metadata, layout, scale 

112 return output 

113 

114 @staticmethod 

115 @once_differentiable 

116 def backward(ctx: Any, grad_output: torch.Tensor) -> tuple: 

117 """Return effective-input gradients without applying norm or sigmoid twice.""" 

118 query, key, value, gate_cumsum, beta, matrix_a = ctx.saved_tensors 

119 gradients = _load_ops().chunk_gated_delta_rule_bwd( 

120 query, key, value, gate_cumsum, beta, matrix_a, grad_output.contiguous(), ctx.scale, 

121 layout=ctx.layout, chunk_size=64, initial_state=None, dht=None, use_exp2=False, 

122 use_qk_l2norm_in_kernel=False, use_beta_sigmoid_in_kernel=False, 

123 use_gate_in_kernel=False, state_v_first=False, **ctx.native_metadata, 

124 ) 

125 if len(gradients) != 8 or any(not isinstance(gradient, torch.Tensor) for gradient in gradients[:5]): 

126 raise RuntimeError("fla-npu GDN backward must return eight slots with all five input gradients.") 

127 dq, dk, dv, dbeta, dgate, *_ = gradients 

128 return dq, dk, dv, dgate, dbeta, None 

129 

130 

131def chunk_gated_delta_rule_ascendc( 

132 query: torch.Tensor, 

133 key: torch.Tensor, 

134 value: torch.Tensor, 

135 g: torch.Tensor, 

136 beta: torch.Tensor, 

137 *, 

138 cu_seqlens: torch.Tensor | Sequence[int] | None = None, 

139) -> tuple[torch.Tensor, None]: 

140 """Run the optional native GDN local core for dense or packed Ulysses. 

141 

142 Args: 

143 query: Normalized or unnormalized BF16 query of shape [B,T,H,128]. 

144 key: BF16 key with the same shape as query. 

145 value: BF16 value with the same shape as query. 

146 g: Effective FP32 log-decay of shape [B,T,H], before chunk cumsum. 

147 beta: Effective BF16 delta coefficient of shape [B,T,H]. 

148 cu_seqlens: Immutable document boundaries supplied as host integers. 

149 

150 Returns: 

151 The output with the same shape as value, and None for final state. 

152 

153 Raises: 

154 ValueError: Inputs violate the explicit native training contract. 

155 RuntimeError: The optional native API or device is unsupported. 

156 """ 

157 if query.ndim != 4: 

158 raise ValueError("GDN inputs must have shape [B,T,H,K].") 

159 if cu_seqlens is not None and query.shape[0] != 1: 

160 raise ValueError("Packed Ascend C GDN requires physical batch size one.") 

161 metadata = _packed_metadata(cu_seqlens, query.shape[1]) 

162 return _AscendCGDN.apply(query, key, value, g, beta, metadata), None