Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / kimi_delta_attention_fla_adapter.py: 0%
98 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"""Lazy adapter for the external FLA Triton-Ascend KDA backend."""
16from __future__ import annotations
18# pylint: disable=forbidden-backend-import
20__all__ = [
21 "FLAKDAStagedOps",
22 "get_fla_kda_staged_ops",
23 "is_fla_triton_kda_available",
24 "run_fla_chunk_kda",
25]
27import importlib
28import importlib.util
29import re
30from dataclasses import dataclass
31from functools import lru_cache
32from typing import Any, Callable, Optional
34import torch
37_MIN_FLA_VERSION = (0, 6, 0)
40@dataclass(frozen=True)
41class FLAKDAStagedOps:
42 """FLA operators required around Hyper's affine state summary."""
44 l2norm_fwd: Callable[..., Any]
45 l2norm_bwd: Callable[..., Any]
46 fused_beta_sigmoid: Callable[..., Any]
47 fused_beta_sigmoid_bwd: Callable[..., Any]
48 kda_gate_chunk_cumsum: Callable[..., Any]
49 kda_gate_bwd: Callable[..., Any]
50 chunk_kda_fwd_intra: Callable[..., Any]
51 recompute_w_u_fwd: Callable[..., Any]
52 chunk_gated_delta_rule_fwd_h: Callable[..., Any]
53 chunk_gla_fwd_o_gk: Callable[..., Any]
54 chunk_kda_bwd_dav: Callable[..., Any]
55 chunk_gated_delta_rule_bwd_dhu: Callable[..., Any]
56 chunk_kda_bwd_wy_dqkg_fused: Callable[..., Any]
57 chunk_kda_bwd_intra: Callable[..., Any]
58 chunk_local_cumsum: Callable[..., Any]
61@dataclass(frozen=True)
62class _FLAKDARuntime:
63 """Resolved external FLA runtime and its required KDA entry points."""
65 version: str
66 chunk_kda: Callable[..., Any]
67 staged: FLAKDAStagedOps
70def _get_attribute(module_name: str, attribute: str) -> Any:
71 """Import one required FLA symbol with a useful compatibility error."""
72 try:
73 module = importlib.import_module(module_name)
74 except ModuleNotFoundError as exc:
75 missing_name = getattr(exc, "name", None)
76 is_fla_module = (
77 missing_name is not None and str(missing_name).startswith("fla.")
78 )
79 if missing_name == module_name or is_fla_module:
80 raise RuntimeError(
81 "The installed FLA package is incompatible with Hyper KDA: "
82 f"missing module {module_name}."
83 ) from exc
84 raise RuntimeError(
85 f"FLA module {module_name} could not load runtime dependency "
86 f"{missing_name!r}."
87 ) from exc
88 except ImportError as exc:
89 raise RuntimeError(
90 f"FLA module {module_name} failed to import; inspect the chained "
91 "exception for the incompatible runtime library."
92 ) from exc
93 try:
94 return getattr(module, attribute)
95 except AttributeError as exc:
96 raise RuntimeError(
97 "The installed FLA package is incompatible with Hyper KDA: "
98 f"missing {module_name}.{attribute}. Install an FLA revision with "
99 "the Triton-Ascend KDA backend."
100 ) from exc
103def _parse_version(version: str) -> tuple[int, int, int]:
104 """Parse the numeric prefix of an FLA semantic version."""
105 match = re.match(r"^(\d+)\.(\d+)\.(\d+)", version)
106 if match is None:
107 raise RuntimeError(f"Unable to parse the installed FLA version {version!r}.")
108 return tuple(int(part) for part in match.groups())
111def _require_triton_ascend() -> None:
112 """Require the external FLA runtime to have a Triton Ascend backend."""
113 try:
114 importlib.import_module("triton")
115 ascend_backend = importlib.util.find_spec("triton.backends.ascend")
116 except (ImportError, ModuleNotFoundError) as exc:
117 raise RuntimeError(
118 "KDA backend='triton' requires Triton-Ascend in the runtime environment."
119 ) from exc
120 if ascend_backend is None:
121 raise RuntimeError(
122 "KDA backend='triton' found Triton, but its Ascend backend is missing."
123 )
126def _require_npu_backends() -> None:
127 """Verify that FLA registered every backend used by the staged P2P path."""
128 backend_specs = (
129 (
130 "fla.ops.kda.backends.triton_ascend",
131 "TritonAscendKDABackend",
132 ),
133 (
134 "fla.ops.common.backends.triton_ascend",
135 "TritonAscendCommonBackend",
136 ),
137 (
138 "fla.ops.gla.backends.triton_ascend",
139 "TritonAscendGLABackend",
140 ),
141 )
142 for module_name, class_name in backend_specs:
143 backend = _get_attribute(module_name, class_name)
144 if not backend.is_available():
145 raise RuntimeError(
146 "KDA backend='triton' requires FLA's Triton-Ascend backend "
147 f"{class_name}, but it is unavailable in this process."
148 )
151@lru_cache(maxsize=1)
152def _require_fla_kda_runtime() -> _FLAKDARuntime:
153 """Resolve and validate the optional external FLA dependency once."""
154 try:
155 fla = importlib.import_module("fla")
156 except ModuleNotFoundError as exc:
157 missing_name = getattr(exc, "name", None)
158 if missing_name != "fla":
159 raise RuntimeError(
160 "The external FLA package is present but cannot load runtime "
161 f"dependency {missing_name!r}."
162 ) from exc
163 raise RuntimeError(
164 "KDA backend='triton' requires the optional flash-linear-attention "
165 "package with its Triton-Ascend KDA backend."
166 ) from exc
167 except ImportError as exc:
168 raise RuntimeError(
169 "The external FLA package is present but failed to import; inspect "
170 "the chained exception for the incompatible runtime library."
171 ) from exc
173 version = getattr(fla, "__version__", None)
174 if not isinstance(version, str):
175 raise RuntimeError("The installed FLA package does not expose __version__.")
176 if _parse_version(version) < _MIN_FLA_VERSION:
177 minimum = ".".join(str(part) for part in _MIN_FLA_VERSION)
178 raise RuntimeError(
179 f"KDA backend='triton' requires FLA >= {minimum}, got {version}."
180 )
182 _require_triton_ascend()
183 _require_npu_backends()
184 staged = FLAKDAStagedOps(
185 l2norm_fwd=_get_attribute("fla.modules.l2norm", "l2norm_fwd"),
186 l2norm_bwd=_get_attribute("fla.modules.l2norm", "l2norm_bwd"),
187 fused_beta_sigmoid=_get_attribute(
188 "fla.ops.common.gate", "fused_beta_sigmoid"
189 ),
190 fused_beta_sigmoid_bwd=_get_attribute(
191 "fla.ops.common.gate", "fused_beta_sigmoid_bwd"
192 ),
193 kda_gate_chunk_cumsum=_get_attribute(
194 "fla.ops.kda.gate", "kda_gate_chunk_cumsum"
195 ),
196 kda_gate_bwd=_get_attribute("fla.ops.kda.gate", "kda_gate_bwd"),
197 chunk_kda_fwd_intra=_get_attribute(
198 "fla.ops.kda.chunk_intra", "chunk_kda_fwd_intra"
199 ),
200 recompute_w_u_fwd=_get_attribute(
201 "fla.ops.kda.wy_fast", "recompute_w_u_fwd"
202 ),
203 chunk_gated_delta_rule_fwd_h=_get_attribute(
204 "fla.ops.common.chunk_delta_h", "chunk_gated_delta_rule_fwd_h"
205 ),
206 chunk_gla_fwd_o_gk=_get_attribute(
207 "fla.ops.gla.chunk", "chunk_gla_fwd_o_gk"
208 ),
209 chunk_kda_bwd_dav=_get_attribute(
210 "fla.ops.kda.chunk_bwd", "chunk_kda_bwd_dAv"
211 ),
212 chunk_gated_delta_rule_bwd_dhu=_get_attribute(
213 "fla.ops.common.chunk_delta_h", "chunk_gated_delta_rule_bwd_dhu"
214 ),
215 chunk_kda_bwd_wy_dqkg_fused=_get_attribute(
216 "fla.ops.kda.chunk_bwd", "chunk_kda_bwd_wy_dqkg_fused"
217 ),
218 chunk_kda_bwd_intra=_get_attribute(
219 "fla.ops.kda.chunk_intra", "chunk_kda_bwd_intra"
220 ),
221 chunk_local_cumsum=_get_attribute(
222 "fla.ops.utils.cumsum", "chunk_local_cumsum"
223 ),
224 )
225 return _FLAKDARuntime(
226 version=version,
227 chunk_kda=_get_attribute("fla.ops.kda", "chunk_kda"),
228 staged=staged,
229 )
232def is_fla_triton_kda_available() -> bool:
233 """Return whether the optional external FLA KDA backend is usable."""
234 try:
235 _require_fla_kda_runtime()
236 except RuntimeError:
237 return False
238 return True
241def get_fla_kda_staged_ops() -> FLAKDAStagedOps:
242 """Return FLA staged operators for Hyper's state-P2P autograd function."""
243 return _require_fla_kda_runtime().staged
246def run_fla_chunk_kda(
247 query: torch.Tensor,
248 key: torch.Tensor,
249 value: torch.Tensor,
250 gate: torch.Tensor,
251 beta: torch.Tensor,
252 *,
253 a_log: torch.Tensor,
254 dt_bias: torch.Tensor,
255 scale: Optional[float] = None,
256 initial_state: Optional[torch.Tensor] = None,
257 output_final_state: bool = False,
258 lower_bound: float = -5.0,
259 chunk_size: int = 64,
260 safe_gate: bool = True,
261) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
262 """Run FLA's public KDA API with the validated Kimi K3 fused contract."""
263 runtime = _require_fla_kda_runtime()
264 return runtime.chunk_kda(
265 q=query,
266 k=key,
267 v=value,
268 g=gate,
269 beta=beta,
270 A_log=a_log,
271 dt_bias=dt_bias,
272 scale=scale,
273 initial_state=initial_state,
274 output_final_state=output_final_state,
275 use_qk_l2norm_in_kernel=True,
276 use_gate_in_kernel=True,
277 use_beta_sigmoid_in_kernel=True,
278 safe_gate=safe_gate,
279 lower_bound=lower_bound,
280 chunk_size=chunk_size,
281 )