Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / gated_delta_net_state_summary.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"""Affine state-summary operations used by GDN State-P2P.""" 

16 

17# pylint: disable=forbidden-backend-import 

18 

19__all__ = [ 

20 "apply_gdn_state_gradient_summary", 

21 "apply_gdn_state_summary", 

22 "chunk_gated_delta_rule_state_gradient_summary_bwd", 

23 "chunk_gated_delta_rule_state_summary_fwd", 

24] 

25 

26from typing import Optional 

27 

28import torch 

29import triton 

30 

31from ._triton.gated_delta_net.state_summary import ( 

32 gdn_packed_state_summary_kernel, 

33 gdn_state_grad_ext_kernel, 

34) 

35 

36 

37def _validate_fixed_summary_shape( 

38 key_dim: int, 

39 value_dim: int, 

40 chunk_size: int, 

41) -> None: 

42 if key_dim != 128 or value_dim != 128 or chunk_size != 64: 

43 raise NotImplementedError( 

44 "Triton GDN state summary requires key_dim=value_dim=128 and " 

45 "chunk_size=64." 

46 ) 

47 

48 

49@torch.compiler.disable 

50def chunk_gated_delta_rule_state_summary_fwd( 

51 key: torch.Tensor, 

52 w: torch.Tensor, 

53 u: torch.Tensor, 

54 g: torch.Tensor, 

55 *, 

56 chunk_size: int = 64, 

57 block_size: int = 128, 

58 output_transition: bool = True, 

59) -> tuple[torch.Tensor, Optional[torch.Tensor]]: 

60 """Return S and optionally M of ``state_out = M @ state_in + S``. 

61 

62 Set output_transition=False only when the caller has no consumer for M. 

63 Both modes retain the same packed kernel; omitting M avoids its output 

64 copy and saved-state lifetime, not its arithmetic. 

65 """ 

66 if (key.ndim, w.ndim, u.ndim, g.ndim) != (4, 4, 4, 3): 

67 raise ValueError("GDN state summary expects key/w/u [B,T,H,D] and g [B,T,H].") 

68 batch, seq_len, heads, key_dim = key.shape 

69 value_dim = u.shape[-1] 

70 _validate_fixed_summary_shape(key_dim, value_dim, chunk_size) 

71 if block_size not in (64, 128): 

72 raise ValueError(f"GDN state-summary block_size must be 64 or 128, got {block_size}.") 

73 if seq_len <= 0: 

74 raise ValueError("State summary requires a nonempty sample fragment.") 

75 if w.shape != key.shape or u.shape[:3] != key.shape[:3] or g.shape != key.shape[:3]: 

76 raise ValueError( 

77 "Incompatible GDN state-summary shapes: " 

78 f"key={tuple(key.shape)}, w={tuple(w.shape)}, " 

79 f"u={tuple(u.shape)}, g={tuple(g.shape)}." 

80 ) 

81 

82 key, w, u, g = (tensor.contiguous() for tensor in (key, w, u, g)) 

83 packed_summary = torch.empty( 

84 batch, 

85 heads, 

86 key_dim, 

87 value_dim + key_dim, 

88 device=key.device, 

89 dtype=torch.float32, 

90 ) 

91 gdn_packed_state_summary_kernel[ 

92 (triton.cdiv(value_dim + key_dim, block_size), batch * heads) 

93 ]( 

94 key, 

95 w, 

96 u, 

97 g, 

98 packed_summary, 

99 seq_len, 

100 H=heads, 

101 K=key_dim, 

102 V=value_dim, 

103 BT=chunk_size, 

104 BV=block_size, 

105 NT=triton.cdiv(seq_len, chunk_size), 

106 ) 

107 state_ext = packed_summary[..., :value_dim].contiguous() 

108 if not output_transition: 

109 return state_ext, None 

110 transition = packed_summary[..., value_dim:].contiguous() 

111 return state_ext, transition 

112 

113 

114def _validate_gdn_gradient_summary_inputs( 

115 query: torch.Tensor, 

116 key: torch.Tensor, 

117 w: torch.Tensor, 

118 g: torch.Tensor, 

119 grad_output: torch.Tensor, 

120 dv: torch.Tensor, 

121 chunk_size: int, 

122) -> tuple[int, int, int, int, int]: 

123 """Validate the reverse summary inputs before launching its kernel.""" 

124 batch, seq_len, heads, key_dim = query.shape 

125 value_dim = grad_output.shape[-1] 

126 _validate_fixed_summary_shape(key_dim, value_dim, chunk_size) 

127 if seq_len <= 0: 

128 raise ValueError("State summary requires a nonempty sample fragment.") 

129 qk_shape = (batch, seq_len, heads, key_dim) 

130 value_shape = (batch, seq_len, heads, value_dim) 

131 if (key.shape, w.shape, g.shape, grad_output.shape, dv.shape) != ( 

132 qk_shape, qk_shape, qk_shape[:3], value_shape, value_shape 

133 ): 

134 raise ValueError( 

135 "Incompatible GDN state-gradient summary shapes: " 

136 f"query={tuple(query.shape)}, key={tuple(key.shape)}, " 

137 f"w={tuple(w.shape)}, g={tuple(g.shape)}, " 

138 f"grad_output={tuple(grad_output.shape)}, dv={tuple(dv.shape)}." 

139 ) 

140 return batch, seq_len, heads, key_dim, value_dim 

141 

142 

143@torch.compiler.disable 

144def chunk_gated_delta_rule_state_gradient_summary_bwd( 

145 query: torch.Tensor, 

146 key: torch.Tensor, 

147 w: torch.Tensor, 

148 g: torch.Tensor, 

149 grad_output: torch.Tensor, 

150 dv: torch.Tensor, 

151 scale: float, 

152 *, 

153 chunk_size: int = 64, 

154) -> torch.Tensor: 

155 """Return the local-loss contribution to the incoming state gradient.""" 

156 batch, seq_len, heads, key_dim, value_dim = _validate_gdn_gradient_summary_inputs( 

157 query, key, w, g, grad_output, dv, chunk_size 

158 ) 

159 

160 query, key, w, g, grad_output, dv = ( 

161 tensor.contiguous() for tensor in (query, key, w, g, grad_output, dv) 

162 ) 

163 grad_state_ext = torch.empty( 

164 batch, 

165 heads, 

166 key_dim, 

167 value_dim, 

168 device=query.device, 

169 dtype=torch.float32, 

170 ) 

171 gdn_state_grad_ext_kernel[(1, batch * heads)]( 

172 query, 

173 key, 

174 w, 

175 g, 

176 grad_output, 

177 dv, 

178 grad_state_ext, 

179 scale, 

180 seq_len, 

181 H=heads, 

182 K=key_dim, 

183 V=value_dim, 

184 BT=chunk_size, 

185 BV=128, 

186 NT=triton.cdiv(seq_len, chunk_size), 

187 ) 

188 return grad_state_ext 

189 

190 

191def apply_gdn_state_summary( 

192 state_ext: torch.Tensor, 

193 transition: Optional[torch.Tensor], 

194 initial_state: Optional[torch.Tensor], 

195) -> torch.Tensor: 

196 """Apply a local affine state summary in FP32.""" 

197 if initial_state is None: 

198 return state_ext 

199 return torch.matmul(transition, initial_state.float()) + state_ext 

200 

201 

202def apply_gdn_state_gradient_summary( 

203 grad_state_ext: torch.Tensor, 

204 transition: Optional[torch.Tensor], 

205 grad_final_state: Optional[torch.Tensor], 

206) -> torch.Tensor: 

207 """Apply the adjoint affine summary to a gradient from the next rank.""" 

208 if grad_final_state is None: 

209 return grad_state_ext 

210 return ( 

211 torch.matmul(transition.transpose(-2, -1), grad_final_state.float()) 

212 + grad_state_ext 

213 )