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) -> tuple[torch.Tensor, torch.Tensor]: 

59 """Return the local affine map ``state_out = M @ state_in + S``.""" 

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

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

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

63 value_dim = u.shape[-1] 

64 _validate_fixed_summary_shape(key_dim, value_dim, chunk_size) 

65 if block_size not in (64, 128): 

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

67 if seq_len % chunk_size != 0: 

68 raise ValueError( 

69 f"GDN state-summary sequence length {seq_len} must be divisible by {chunk_size}." 

70 ) 

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

72 raise ValueError( 

73 "Incompatible GDN state-summary shapes: " 

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

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

76 ) 

77 

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

79 packed_summary = torch.empty( 

80 batch, 

81 heads, 

82 key_dim, 

83 value_dim + key_dim, 

84 device=key.device, 

85 dtype=torch.float32, 

86 ) 

87 gdn_packed_state_summary_kernel[ 

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

89 ]( 

90 key, 

91 w, 

92 u, 

93 g, 

94 packed_summary, 

95 seq_len, 

96 H=heads, 

97 K=key_dim, 

98 V=value_dim, 

99 BT=chunk_size, 

100 BV=block_size, 

101 NT=seq_len // chunk_size, 

102 ) 

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

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

105 return state_ext, transition 

106 

107 

108def _validate_gdn_gradient_summary_inputs( 

109 query: torch.Tensor, 

110 key: torch.Tensor, 

111 w: torch.Tensor, 

112 g: torch.Tensor, 

113 grad_output: torch.Tensor, 

114 dv: torch.Tensor, 

115 chunk_size: int, 

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

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

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

119 value_dim = grad_output.shape[-1] 

120 _validate_fixed_summary_shape(key_dim, value_dim, chunk_size) 

121 if seq_len % chunk_size != 0: 

122 raise ValueError( 

123 f"GDN state-gradient sequence length {seq_len} must be divisible by {chunk_size}." 

124 ) 

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

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

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

128 qk_shape, qk_shape, qk_shape[:3], value_shape, value_shape 

129 ): 

130 raise ValueError( 

131 "Incompatible GDN state-gradient summary shapes: " 

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

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

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

135 ) 

136 return batch, seq_len, heads, key_dim, value_dim 

137 

138 

139@torch.compiler.disable 

140def chunk_gated_delta_rule_state_gradient_summary_bwd( 

141 query: torch.Tensor, 

142 key: torch.Tensor, 

143 w: torch.Tensor, 

144 g: torch.Tensor, 

145 grad_output: torch.Tensor, 

146 dv: torch.Tensor, 

147 scale: float, 

148 *, 

149 chunk_size: int = 64, 

150) -> torch.Tensor: 

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

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

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

154 ) 

155 

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

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

158 ) 

159 grad_state_ext = torch.empty( 

160 batch, 

161 heads, 

162 key_dim, 

163 value_dim, 

164 device=query.device, 

165 dtype=torch.float32, 

166 ) 

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

168 query, 

169 key, 

170 w, 

171 g, 

172 grad_output, 

173 dv, 

174 grad_state_ext, 

175 scale, 

176 seq_len, 

177 H=heads, 

178 K=key_dim, 

179 V=value_dim, 

180 BT=chunk_size, 

181 BV=128, 

182 NT=seq_len // chunk_size, 

183 ) 

184 return grad_state_ext 

185 

186 

187def apply_gdn_state_summary( 

188 state_ext: torch.Tensor, 

189 transition: torch.Tensor, 

190 initial_state: Optional[torch.Tensor], 

191) -> torch.Tensor: 

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

193 if initial_state is None: 

194 return state_ext 

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

196 

197 

198def apply_gdn_state_gradient_summary( 

199 grad_state_ext: torch.Tensor, 

200 transition: torch.Tensor, 

201 grad_final_state: Optional[torch.Tensor], 

202) -> torch.Tensor: 

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

204 if grad_final_state is None: 

205 return grad_state_ext 

206 return ( 

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

208 + grad_state_ext 

209 )