Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / gated_delta_net.py: 64%
117 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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"""Gated DeltaNet eager oracle and explicit local-backend dispatcher."""
16# pylint: disable=forbidden-backend-import,missing-public-type-hints
17# pylint: disable=missing-public-docstring,not-callable
19__all__ = ["chunk_gated_delta_rule", "is_triton_gdn_available", "torch_chunk_gated_delta_rule"]
21import importlib
22import importlib.metadata
23import importlib.util
24import re
25from typing import Optional
27import torch
28from torch.nn import functional as F
30from hyper_parallel.components.functional.linear_attention_packed import (
31 packed_reference,
32)
34_GDN_BACKENDS = frozenset({"eager", "triton", "ascendc"})
35_MIN_TRITON_ASCEND_VERSION = (3, 2, 1)
36_TRITON_GDN_HEAD_DIM = 128
37_TRITON_GDN_CHUNK_SIZE = 64
40def _l2norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor:
41 """L2-normalize along ``dim``."""
42 # ``(x * x).sum`` (MulBackward) instead of ``x.pow(2).sum`` (PowBackward):
43 # equal math, NPU yields ULP-different gradients across the two ops.
44 return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
47def torch_chunk_gated_delta_rule(
48 query: torch.Tensor,
49 key: torch.Tensor,
50 value: torch.Tensor,
51 g: torch.Tensor,
52 beta: torch.Tensor,
53 chunk_size: int = 64,
54 initial_state: Optional[torch.Tensor] = None,
55 output_final_state: bool = False,
56 use_qk_l2norm_in_kernel: bool = False,
57 cu_seqlens: Optional[torch.Tensor] = None,
58):
59 """Chunked gated delta-rule linear attention (pure torch reference).
61 fp32 internal compute, output cast back to input dtype.
63 Shapes
64 ------
65 query / key: ``(B, S, num_k_heads, head_k_dim)``
66 value: ``(B, S, num_v_heads, head_v_dim)`` (here num_k_heads == num_v_heads)
67 g: ``(B, S, num_v_heads)``
68 beta: ``(B, S, num_v_heads)``
69 """
70 if cu_seqlens is not None:
71 return packed_reference(
72 torch_chunk_gated_delta_rule, (query, key, value, g, beta), cu_seqlens,
73 initial_state=initial_state, output_final_state=output_final_state,
74 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, chunk_size=chunk_size,
75 )
76 initial_dtype = query.dtype
77 if use_qk_l2norm_in_kernel:
78 query = _l2norm(query, dim=-1, eps=1e-6)
79 key = _l2norm(key, dim=-1, eps=1e-6)
80 # (B, S, H, D) → (B, H, S, D), fp32
81 query, key, value, beta, g = [
82 x.transpose(1, 2).contiguous().to(torch.float32)
83 for x in (query, key, value, beta, g)
84 ]
86 batch_size, num_heads, sequence_length, k_head_dim = key.shape
87 v_head_dim = value.shape[-1]
88 pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
89 query = F.pad(query, (0, 0, 0, pad_size))
90 key = F.pad(key, (0, 0, 0, pad_size))
91 value = F.pad(value, (0, 0, 0, pad_size))
92 beta = F.pad(beta, (0, pad_size))
93 g = F.pad(g, (0, pad_size))
94 total_sequence_length = sequence_length + pad_size
95 scale = 1 / (query.shape[-1] ** 0.5)
96 query = query * scale
98 v_beta = value * beta.unsqueeze(-1)
99 k_beta = key * beta.unsqueeze(-1)
100 # (B, H, S, D) → (B, H, n_chunks, chunk_size, D)
101 query, key, value, k_beta, v_beta = [
102 x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
103 for x in (query, key, value, k_beta, v_beta)
104 ]
105 g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
106 mask = torch.triu(
107 torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
108 diagonal=0,
109 )
111 # chunk decay
112 g = g.cumsum(dim=-1)
113 decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
114 attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
115 for i in range(1, chunk_size):
116 row = attn[..., i, :i].clone()
117 sub = attn[..., :i, :i].clone()
118 attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
119 attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
120 value = attn @ v_beta
121 k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
122 last_recurrent_state = (
123 torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
124 if initial_state is None
125 else initial_state.to(value)
126 )
127 core_attn_out = torch.zeros_like(value)
129 # for each chunk
130 for i in range(0, total_sequence_length // chunk_size):
131 q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
132 attn = q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]
133 v_prime = k_cumdecay[:, :, i] @ last_recurrent_state
134 v_new = v_i - v_prime
135 attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state
136 core_attn_out[:, :, i] = attn_inter + attn @ v_new
137 last_recurrent_state = (
138 last_recurrent_state * g[:, :, i, -1, None, None].exp()
139 + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new
140 )
142 if not output_final_state:
143 last_recurrent_state = None
144 core_attn_out = core_attn_out.reshape(
145 core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1],
146 )
147 core_attn_out = core_attn_out[:, :, :sequence_length]
148 core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
149 return core_attn_out, last_recurrent_state
152def _parse_version(version_text: str) -> tuple[int, int, int]:
153 """Return a three-component numeric version tuple."""
154 parts = []
155 for part in version_text.split("+")[0].split(".")[:3]:
156 match = re.match(r"\d+", part)
157 if match is not None:
158 parts.append(int(match.group()))
159 return tuple((parts + [0, 0, 0])[:3])
162def _is_triton_gdn_input_supported(
163 query: torch.Tensor,
164 key: Optional[torch.Tensor],
165 value: Optional[torch.Tensor],
166 g: Optional[torch.Tensor],
167 beta: Optional[torch.Tensor],
168) -> bool:
169 """Check the fixed Qwen3.5 GDN contract validated by this backend."""
170 if any(tensor is None for tensor in (key, value, g, beta)):
171 return False
172 if not (
173 query.device.type == "npu"
174 and query.dtype == key.dtype == value.dtype == beta.dtype == torch.bfloat16
175 and g.dtype == torch.float32
176 and query.ndim == key.ndim == value.ndim == 4
177 and g.ndim == beta.ndim == 3
178 and query.shape == key.shape
179 and query.shape[:3] == value.shape[:3] == g.shape == beta.shape
180 and query.shape[-1] == value.shape[-1] == _TRITON_GDN_HEAD_DIM
181 ):
182 return False
183 return all(
184 tensor.device == query.device
185 for tensor in (key, value, g, beta)
186 )
189def is_triton_gdn_available(
190 query: Optional[torch.Tensor] = None,
191 key: Optional[torch.Tensor] = None,
192 value: Optional[torch.Tensor] = None,
193 g: Optional[torch.Tensor] = None,
194 beta: Optional[torch.Tensor] = None,
195 chunk_size: int = _TRITON_GDN_CHUNK_SIZE,
196) -> bool:
197 """Return whether the validated Triton-Ascend GDN backend is available."""
198 if chunk_size != _TRITON_GDN_CHUNK_SIZE:
199 return False
200 if query is not None and not _is_triton_gdn_input_supported(
201 query, key, value, g, beta,
202 ):
203 return False
204 try:
205 version_text = importlib.metadata.version("triton-ascend")
206 except importlib.metadata.PackageNotFoundError:
207 return False
208 version = _parse_version(version_text)
209 if version < _MIN_TRITON_ASCEND_VERSION:
210 return False
211 try:
212 importlib.import_module("triton")
213 return importlib.util.find_spec("triton.backends.ascend") is not None
214 except (AttributeError, ImportError, ModuleNotFoundError):
215 return False
218def chunk_gated_delta_rule(
219 query: torch.Tensor,
220 key: torch.Tensor,
221 value: torch.Tensor,
222 g: torch.Tensor,
223 beta: torch.Tensor,
224 chunk_size: int = 64,
225 initial_state: Optional[torch.Tensor] = None,
226 output_final_state: bool = False,
227 use_qk_l2norm_in_kernel: bool = False,
228 backend: str = "eager",
229 cu_seqlens: Optional[torch.Tensor] = None,
230):
231 """Dispatch GDN to an explicitly selected eager, Triton or Ascend C backend."""
232 if query.ndim == 3:
233 output, final = chunk_gated_delta_rule(
234 *(tensor.unsqueeze(0) for tensor in (query, key, value, g, beta)),
235 initial_state=initial_state, output_final_state=output_final_state,
236 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, chunk_size=chunk_size,
237 backend=backend, cu_seqlens=cu_seqlens)
238 return output.squeeze(0), final
239 backend = backend.lower()
240 if backend not in _GDN_BACKENDS:
241 raise ValueError(
242 f"unsupported GDN backend {backend!r}; "
243 f"expected one of {sorted(_GDN_BACKENDS)}."
244 )
245 if backend == "eager":
246 return torch_chunk_gated_delta_rule(
247 query,
248 key,
249 value,
250 g=g,
251 beta=beta,
252 chunk_size=chunk_size,
253 initial_state=initial_state,
254 output_final_state=output_final_state,
255 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
256 cu_seqlens=cu_seqlens,
257 )
258 if backend == "ascendc":
259 if initial_state is not None or output_final_state:
260 raise NotImplementedError("GDN Ascend C currently supports zero-state local or pure Ulysses training only.")
261 if chunk_size != 64:
262 raise ValueError("GDN Ascend C requires chunk_size=64.")
263 # Preserve the existing GDN normalization graph outside the native core.
264 if use_qk_l2norm_in_kernel:
265 query, key = _l2norm(query), _l2norm(key)
266 # Keep optional native dependencies out of eager and Triton imports.
267 from hyper_parallel.components.functional.gated_delta_net_ascendc import ( # pylint: disable=import-outside-toplevel
268 chunk_gated_delta_rule_ascendc,
269 )
270 return chunk_gated_delta_rule_ascendc(query, key, value, g, beta, cu_seqlens=cu_seqlens)
271 if not is_triton_gdn_available(
272 query, key, value, g, beta, chunk_size=chunk_size,
273 ):
274 raise RuntimeError(
275 "GDN backend='triton' requires triton-ascend >= 3.2.1 "
276 "installation with the Ascend backend and NPU inputs "
277 "q/k/v/beta=bf16, g=fp32, head_k_dim=head_v_dim=128, "
278 "and chunk_size=64."
279 )
281 from hyper_parallel.components.functional.gated_delta_net import ( # pylint: disable=import-outside-toplevel
282 chunk_gated_delta_rule as triton_chunk_gated_delta_rule,
283 )
285 return triton_chunk_gated_delta_rule(
286 query,
287 key,
288 value,
289 g,
290 beta,
291 initial_state=initial_state,
292 output_final_state=output_final_state,
293 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
294 chunk_size=chunk_size,
295 cu_seqlens=cu_seqlens,
296 )