Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / custom_ops / gdn / state_summary.py: 0%
50 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +0800
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."""
17from typing import Optional
19import torch
20import triton
22from .triton.state_summary import (
23 gdn_packed_state_summary_kernel,
24 gdn_state_grad_ext_kernel,
25)
28def _validate_fixed_summary_shape(
29 key_dim: int,
30 value_dim: int,
31 chunk_size: int,
32) -> None:
33 if key_dim != 128 or value_dim != 128 or chunk_size != 64:
34 raise NotImplementedError(
35 "Triton GDN state summary requires key_dim=value_dim=128 and "
36 "chunk_size=64."
37 )
40@torch.compiler.disable
41def chunk_gated_delta_rule_state_summary_fwd(
42 key: torch.Tensor,
43 w: torch.Tensor,
44 u: torch.Tensor,
45 g: torch.Tensor,
46 *,
47 chunk_size: int = 64,
48 block_size: int = 128,
49) -> tuple[torch.Tensor, torch.Tensor]:
50 """Return the local affine map ``state_out = M @ state_in + S``."""
51 if key.ndim != 4 or w.ndim != 4 or u.ndim != 4 or g.ndim != 3:
52 raise ValueError("GDN state summary expects key/w/u [B,T,H,D] and g [B,T,H].")
53 batch, seq_len, heads, key_dim = key.shape
54 value_dim = u.shape[-1]
55 _validate_fixed_summary_shape(key_dim, value_dim, chunk_size)
56 if block_size not in (64, 128):
57 raise ValueError(f"GDN state-summary block_size must be 64 or 128, got {block_size}.")
58 if seq_len % chunk_size != 0:
59 raise ValueError(
60 f"GDN state-summary sequence length {seq_len} must be divisible by {chunk_size}."
61 )
62 if w.shape != key.shape or u.shape[:3] != key.shape[:3] or g.shape != key.shape[:3]:
63 raise ValueError(
64 "Incompatible GDN state-summary shapes: "
65 f"key={tuple(key.shape)}, w={tuple(w.shape)}, "
66 f"u={tuple(u.shape)}, g={tuple(g.shape)}."
67 )
69 key, w, u, g = (tensor.contiguous() for tensor in (key, w, u, g))
70 packed_summary = torch.empty(
71 batch,
72 heads,
73 key_dim,
74 value_dim + key_dim,
75 device=key.device,
76 dtype=torch.float32,
77 )
78 gdn_packed_state_summary_kernel[
79 (triton.cdiv(value_dim + key_dim, block_size), batch * heads)
80 ](
81 key,
82 w,
83 u,
84 g,
85 packed_summary,
86 seq_len,
87 H=heads,
88 K=key_dim,
89 V=value_dim,
90 BT=chunk_size,
91 BV=block_size,
92 NT=seq_len // chunk_size,
93 )
94 state_ext = packed_summary[..., :value_dim].contiguous()
95 transition = packed_summary[..., value_dim:].contiguous()
96 return state_ext, transition
99@torch.compiler.disable
100def chunk_gated_delta_rule_state_gradient_summary_bwd(
101 query: torch.Tensor,
102 key: torch.Tensor,
103 w: torch.Tensor,
104 g: torch.Tensor,
105 grad_output: torch.Tensor,
106 dv: torch.Tensor,
107 scale: float,
108 *,
109 chunk_size: int = 64,
110) -> torch.Tensor:
111 """Return the local-loss contribution to the incoming state gradient."""
112 batch, seq_len, heads, key_dim = query.shape
113 value_dim = grad_output.shape[-1]
114 _validate_fixed_summary_shape(key_dim, value_dim, chunk_size)
115 if seq_len % chunk_size != 0:
116 raise ValueError(
117 f"GDN state-gradient sequence length {seq_len} must be divisible by {chunk_size}."
118 )
119 qk_shape = (batch, seq_len, heads, key_dim)
120 value_shape = (batch, seq_len, heads, value_dim)
121 if (
122 key.shape != qk_shape
123 or w.shape != qk_shape
124 or g.shape != qk_shape[:3]
125 or grad_output.shape != value_shape
126 or dv.shape != value_shape
127 ):
128 raise ValueError(
129 "Incompatible GDN state-gradient summary shapes: "
130 f"query={tuple(query.shape)}, key={tuple(key.shape)}, "
131 f"w={tuple(w.shape)}, g={tuple(g.shape)}, "
132 f"grad_output={tuple(grad_output.shape)}, dv={tuple(dv.shape)}."
133 )
135 query, key, w, g, grad_output, dv = (
136 tensor.contiguous() for tensor in (query, key, w, g, grad_output, dv)
137 )
138 grad_state_ext = torch.empty(
139 batch,
140 heads,
141 key_dim,
142 value_dim,
143 device=query.device,
144 dtype=torch.float32,
145 )
146 gdn_state_grad_ext_kernel[(1, batch * heads)](
147 query,
148 key,
149 w,
150 g,
151 grad_output,
152 dv,
153 grad_state_ext,
154 scale,
155 seq_len,
156 H=heads,
157 K=key_dim,
158 V=value_dim,
159 BT=chunk_size,
160 BV=128,
161 NT=seq_len // chunk_size,
162 )
163 return grad_state_ext
166def apply_gdn_state_summary(
167 state_ext: torch.Tensor,
168 transition: torch.Tensor,
169 initial_state: Optional[torch.Tensor],
170) -> torch.Tensor:
171 """Apply a local affine state summary in FP32."""
172 if initial_state is None:
173 return state_ext
174 return torch.matmul(transition, initial_state.float()) + state_ext
177def apply_gdn_state_gradient_summary(
178 grad_state_ext: torch.Tensor,
179 transition: torch.Tensor,
180 grad_final_state: Optional[torch.Tensor],
181) -> torch.Tensor:
182 """Apply the adjoint affine summary to a gradient from the next rank."""
183 if grad_final_state is None:
184 return grad_state_ext
185 return (
186 torch.matmul(transition.transpose(-2, -1), grad_final_state.float())
187 + grad_state_ext
188 )
191__all__ = [
192 "apply_gdn_state_gradient_summary",
193 "apply_gdn_state_summary",
194 "chunk_gated_delta_rule_state_gradient_summary_bwd",
195 "chunk_gated_delta_rule_state_summary_fwd",
196]