Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / gated_delta_net.py: 61%
102 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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
30_GDN_BACKENDS = frozenset({"eager", "triton"})
31_MIN_TRITON_ASCEND_VERSION = (3, 2, 1)
32_TRITON_GDN_HEAD_DIM = 128
33_TRITON_GDN_CHUNK_SIZE = 64
36def _l2norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor:
37 """L2-normalize along ``dim``."""
38 # ``(x * x).sum`` (MulBackward) instead of ``x.pow(2).sum`` (PowBackward):
39 # equal math, NPU yields ULP-different gradients across the two ops.
40 return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
43def torch_chunk_gated_delta_rule(
44 query: torch.Tensor,
45 key: torch.Tensor,
46 value: torch.Tensor,
47 g: torch.Tensor,
48 beta: torch.Tensor,
49 chunk_size: int = 64,
50 initial_state: Optional[torch.Tensor] = None,
51 output_final_state: bool = False,
52 use_qk_l2norm_in_kernel: bool = False,
53):
54 """Chunked gated delta-rule linear attention (pure torch reference).
56 fp32 internal compute, output cast back to input dtype.
58 Shapes
59 ------
60 query / key: ``(B, S, num_k_heads, head_k_dim)``
61 value: ``(B, S, num_v_heads, head_v_dim)`` (here num_k_heads == num_v_heads)
62 g: ``(B, S, num_v_heads)``
63 beta: ``(B, S, num_v_heads)``
64 """
65 initial_dtype = query.dtype
66 if use_qk_l2norm_in_kernel:
67 query = _l2norm(query, dim=-1, eps=1e-6)
68 key = _l2norm(key, dim=-1, eps=1e-6)
69 # (B, S, H, D) → (B, H, S, D), fp32
70 query, key, value, beta, g = [
71 x.transpose(1, 2).contiguous().to(torch.float32)
72 for x in (query, key, value, beta, g)
73 ]
75 batch_size, num_heads, sequence_length, k_head_dim = key.shape
76 v_head_dim = value.shape[-1]
77 pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
78 query = F.pad(query, (0, 0, 0, pad_size))
79 key = F.pad(key, (0, 0, 0, pad_size))
80 value = F.pad(value, (0, 0, 0, pad_size))
81 beta = F.pad(beta, (0, pad_size))
82 g = F.pad(g, (0, pad_size))
83 total_sequence_length = sequence_length + pad_size
84 scale = 1 / (query.shape[-1] ** 0.5)
85 query = query * scale
87 v_beta = value * beta.unsqueeze(-1)
88 k_beta = key * beta.unsqueeze(-1)
89 # (B, H, S, D) → (B, H, n_chunks, chunk_size, D)
90 query, key, value, k_beta, v_beta = [
91 x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
92 for x in (query, key, value, k_beta, v_beta)
93 ]
94 g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
95 mask = torch.triu(
96 torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
97 diagonal=0,
98 )
100 # chunk decay
101 g = g.cumsum(dim=-1)
102 decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
103 attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
104 for i in range(1, chunk_size):
105 row = attn[..., i, :i].clone()
106 sub = attn[..., :i, :i].clone()
107 attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
108 attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
109 value = attn @ v_beta
110 k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
111 last_recurrent_state = (
112 torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
113 if initial_state is None
114 else initial_state.to(value)
115 )
116 core_attn_out = torch.zeros_like(value)
118 # for each chunk
119 for i in range(0, total_sequence_length // chunk_size):
120 q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
121 attn = q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]
122 v_prime = k_cumdecay[:, :, i] @ last_recurrent_state
123 v_new = v_i - v_prime
124 attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state
125 core_attn_out[:, :, i] = attn_inter + attn @ v_new
126 last_recurrent_state = (
127 last_recurrent_state * g[:, :, i, -1, None, None].exp()
128 + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new
129 )
131 if not output_final_state:
132 last_recurrent_state = None
133 core_attn_out = core_attn_out.reshape(
134 core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1],
135 )
136 core_attn_out = core_attn_out[:, :, :sequence_length]
137 core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
138 return core_attn_out, last_recurrent_state
141def _parse_version(version_text: str) -> tuple[int, int, int]:
142 """Return a three-component numeric version tuple."""
143 parts = []
144 for part in version_text.split("+")[0].split(".")[:3]:
145 match = re.match(r"\d+", part)
146 if match is not None:
147 parts.append(int(match.group()))
148 return tuple((parts + [0, 0, 0])[:3])
151def _is_triton_gdn_input_supported(
152 query: torch.Tensor,
153 key: Optional[torch.Tensor],
154 value: Optional[torch.Tensor],
155 g: Optional[torch.Tensor],
156 beta: Optional[torch.Tensor],
157) -> bool:
158 """Check the fixed Qwen3.5 GDN contract validated by this backend."""
159 if any(tensor is None for tensor in (key, value, g, beta)):
160 return False
161 if not (
162 query.device.type == "npu"
163 and query.dtype == key.dtype == value.dtype == beta.dtype == torch.bfloat16
164 and g.dtype == torch.float32
165 and query.ndim == key.ndim == value.ndim == 4
166 and g.ndim == beta.ndim == 3
167 and query.shape == key.shape
168 and query.shape[:3] == value.shape[:3] == g.shape == beta.shape
169 and query.shape[-1] == value.shape[-1] == _TRITON_GDN_HEAD_DIM
170 and query.shape[1] % _TRITON_GDN_CHUNK_SIZE == 0
171 ):
172 return False
173 return all(
174 tensor.device == query.device
175 for tensor in (key, value, g, beta)
176 )
179def is_triton_gdn_available(
180 query: Optional[torch.Tensor] = None,
181 key: Optional[torch.Tensor] = None,
182 value: Optional[torch.Tensor] = None,
183 g: Optional[torch.Tensor] = None,
184 beta: Optional[torch.Tensor] = None,
185 chunk_size: int = _TRITON_GDN_CHUNK_SIZE,
186) -> bool:
187 """Return whether the validated Triton-Ascend GDN backend is available."""
188 if chunk_size != _TRITON_GDN_CHUNK_SIZE:
189 return False
190 if query is not None and not _is_triton_gdn_input_supported(
191 query, key, value, g, beta
192 ):
193 return False
194 try:
195 version_text = importlib.metadata.version("triton-ascend")
196 except importlib.metadata.PackageNotFoundError:
197 return False
198 version = _parse_version(version_text)
199 if version < _MIN_TRITON_ASCEND_VERSION:
200 return False
201 try:
202 importlib.import_module("triton")
203 return importlib.util.find_spec("triton.backends.ascend") is not None
204 except (AttributeError, ImportError, ModuleNotFoundError):
205 return False
208def chunk_gated_delta_rule(
209 query: torch.Tensor,
210 key: torch.Tensor,
211 value: torch.Tensor,
212 g: torch.Tensor,
213 beta: torch.Tensor,
214 chunk_size: int = 64,
215 initial_state: Optional[torch.Tensor] = None,
216 output_final_state: bool = False,
217 use_qk_l2norm_in_kernel: bool = False,
218 backend: str = "eager",
219):
220 """Dispatch GDN to an explicitly selected eager or Triton backend."""
221 backend = backend.lower()
222 if backend not in _GDN_BACKENDS:
223 raise ValueError(
224 f"unsupported GDN backend {backend!r}; "
225 f"expected one of {sorted(_GDN_BACKENDS)}."
226 )
227 if backend == "eager":
228 return torch_chunk_gated_delta_rule(
229 query,
230 key,
231 value,
232 g=g,
233 beta=beta,
234 chunk_size=chunk_size,
235 initial_state=initial_state,
236 output_final_state=output_final_state,
237 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
238 )
239 if not is_triton_gdn_available(
240 query, key, value, g, beta, chunk_size=chunk_size
241 ):
242 raise RuntimeError(
243 "GDN backend='triton' requires triton-ascend >= 3.2.1 "
244 "installation with the Ascend backend and NPU inputs "
245 "q/k/v/beta=bf16, g=fp32, head_k_dim=head_v_dim=128, "
246 "chunk_size=64, and sequence length divisible by 64."
247 )
249 from hyper_parallel.components.functional.gated_delta_net import ( # pylint: disable=import-outside-toplevel
250 chunk_gated_delta_rule as triton_chunk_gated_delta_rule,
251 )
253 return triton_chunk_gated_delta_rule(
254 query,
255 key,
256 value,
257 g,
258 beta,
259 initial_state=initial_state,
260 output_final_state=output_final_state,
261 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
262 chunk_size=chunk_size,
263 )