Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / custom_ops / experimental / experimental_ops.py: 100%
17 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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"""Experimental custom operator implementations.
17Each function is a thin delegation wrapper around ``_platform.custom_ops``,
18which routes to the platform-specific Ascend NPU custom C++ kernel.
19"""
20from typing import Optional, Tuple
22from mindspore import Tensor
24from hyper_parallel.platform import get_platform
26_platform = get_platform()
28_MAX_INT64 = 9223372036854775807
31def npu_dense_lightning_indexer_softmax_lse(
32 query_index,
33 key_index,
34 weights,
35 *,
36 actual_seq_qlen: Optional[Tensor] = None,
37 actual_seq_klen: Optional[Tensor] = None,
38 layout: str = 'BSND',
39 sparse_mode: int = 3,
40 pre_tokens: int = _MAX_INT64,
41 next_tokens: int = _MAX_INT64,
42) -> Tuple:
43 """Compute softmax max/sum indices for Lightning Indexer attention.
45 .. warning::
46 This is an experimental API that subject to change or deletion.
48 Pre-computes the Softmax max and sum values to reduce memory usage.
50 The call is routed through the platform ``custom_ops`` layer, which
51 delegates to a ``DFunction`` wrapping the Ascend custom C++ kernel.
52 DTensor inputs are transparently handled by distributed dispatch.
54 Args:
55 query_index: Lightning Indexer query input (Q̃). dtype bfloat16/float16.
56 key_index: Lightning Indexer key input (K̃). Same dtype as query_index.
57 weights: Weight coefficient (W). dtype bfloat16/float16/float32.
58 actual_seq_qlen: Cumulative query sequence lengths (int32 Tensor).
59 actual_seq_klen: Cumulative key sequence lengths (int32 Tensor).
60 layout: Data layout format — 'BSND' (default) or 'TND'.
61 sparse_mode: Sparse computation mode; only mode 3 is supported.
62 pre_tokens: Preceding token window size for sparse attention (int64).
63 next_tokens: Following token window size for sparse attention (int64).
65 Returns:
66 tuple[Tensor, Tensor]: ``(softmax_max_index, softmax_sum_index)``.
67 """
68 return _platform.custom_ops.npu_dense_lightning_indexer_softmax_lse(
69 query_index, key_index, weights,
70 actual_seq_qlen, actual_seq_klen,
71 layout, sparse_mode, pre_tokens, next_tokens,
72 )
75def npu_dense_lightning_indexer_grad_kl_loss(
76 query,
77 key,
78 query_index,
79 key_index,
80 weights,
81 softmax_max,
82 softmax_sum,
83 softmax_max_index,
84 softmax_sum_index,
85 scale_value,
86 *,
87 query_rope=None,
88 key_rope=None,
89 actual_seq_qlen: Optional[Tensor] = None,
90 actual_seq_klen: Optional[Tensor] = None,
91 layout: str = 'BSND',
92 sparse_mode: int = 3,
93 pre_tokens: int = _MAX_INT64,
94 next_tokens: int = _MAX_INT64,
95) -> Tuple:
96 """Compute backward gradients and KL-divergence loss for dense Lightning Indexer.
98 .. warning::
99 This is an experimental API that subject to change or deletion.
101 The call is routed through the platform ``custom_ops`` layer.
103 Returns:
104 tuple[Tensor, Tensor, Tensor, Tensor]:
105 ``(d_query_index, d_key_index, d_weights, loss)``.
106 """
107 return _platform.custom_ops.npu_dense_lightning_indexer_grad_kl_loss(
108 query, key, query_index, key_index, weights,
109 softmax_max, softmax_sum, softmax_max_index, softmax_sum_index,
110 scale_value,
111 query_rope, key_rope,
112 actual_seq_qlen, actual_seq_klen,
113 layout, sparse_mode,
114 pre_tokens, next_tokens,
115 )
118def npu_sparse_lightning_indexer_grad_kl_loss(
119 query,
120 key,
121 query_index,
122 key_index,
123 weights,
124 sparse_indices,
125 softmax_max,
126 softmax_sum,
127 scale_value,
128 *,
129 query_rope=None,
130 key_rope=None,
131 actual_seq_qlen: Optional[Tensor] = None,
132 actual_seq_klen: Optional[Tensor] = None,
133 layout: str = 'BSND',
134 sparse_mode: int = 3,
135 pre_tokens: int = _MAX_INT64,
136 next_tokens: int = _MAX_INT64,
137) -> Tuple:
138 """Compute backward gradients and KL-divergence loss for sparse Lightning Indexer.
140 .. warning::
141 This is an experimental API that subject to change or deletion.
143 Returns:
144 tuple[Tensor, Tensor, Tensor, Tensor]:
145 ``(d_query_index, d_key_index, d_weights, loss)``.
146 """
147 return _platform.custom_ops.npu_sparse_lightning_indexer_grad_kl_loss(
148 query, key, query_index, key_index, weights,
149 sparse_indices, softmax_max, softmax_sum, scale_value,
150 query_rope, key_rope,
151 actual_seq_qlen, actual_seq_klen,
152 layout, sparse_mode,
153 pre_tokens, next_tokens,
154 )
157def npu_mhc_post(x, h_res, h_out, h_post) -> Tuple:
158 """MHC post-processing with residual connection.
160 .. warning::
161 This is an experimental API that subject to change or deletion.
163 Returns:
164 Tensor: Output tensor with same shape and dtype as x.
165 """
166 return _platform.custom_ops.npu_mhc_post(x, h_res, h_out, h_post)
169def npu_mhc_pre_sinkhorn(
170 x,
171 phi,
172 alpha,
173 bias,
174 *,
175 hc_mult: int = 4,
176 num_iters: int = 20,
177 hc_eps: float = 1e-6,
178 norm_eps: float = 1e-6,
179 out_flag: bool = True,
180) -> Tuple:
181 """MHC pre-processing with Sinkhorn normalization.
183 .. warning::
184 This is an experimental API that subject to change or deletion.
186 Returns:
187 tuple: 8 output tensors.
188 """
189 return _platform.custom_ops.npu_mhc_pre_sinkhorn(
190 x, phi, alpha, bias,
191 hc_mult, num_iters,
192 hc_eps, norm_eps, out_flag,
193 )
196def npu_mhc_pre_clamp_sinkhorn(
197 x,
198 phi,
199 alpha,
200 bias,
201 *,
202 hc_mult: int = 4,
203 num_iters: int = 20,
204 hc_eps: float = 1e-6,
205 norm_eps: float = 1e-6,
206 out_flag: bool = True,
207 clamp_min: float = 0.0,
208 clamp_max: float = 0.0,
209) -> Tuple:
210 """MHC pre-processing with clamp and Sinkhorn normalization.
212 .. warning::
213 This is an experimental API that subject to change or deletion.
215 Returns:
216 tuple: 9 output tensors.
217 """
218 return _platform.custom_ops.npu_mhc_pre_clamp_sinkhorn(
219 x, phi, alpha, bias,
220 hc_mult, num_iters,
221 hc_eps, norm_eps, out_flag,
222 clamp_min, clamp_max,
223 )