Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / custom_ops / experimental / experimental_ops.py: 85%
26 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05: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"""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 )
226def npu_lightning_indexer(
227 query,
228 key,
229 weights,
230 sparse_count: int,
231 *,
232 cu_seq_lens_q: Optional[Tensor] = None,
233 cu_seq_lens_k: Optional[Tensor] = None,
234 cmp_residual_k: Optional[Tensor] = None,
235 block_table: Optional[Tensor] = None,
236 layout: str = 'BSND',
237 sparse_mode: int = 0,
238 cmp_ratio: int = 1,
239 return_value: bool = False,
240) -> Tuple:
241 """Sparse attention preprocessing — select top-K key tokens per query token.
243 .. warning::
244 This is an experimental API that subject to change or deletion.
246 Aligned with the ``lightning_indexer`` benchmark signature: positional
247 ``(query, key, weights, sparse_count)`` (``sparse_count`` is benchmark
248 ``topk``); the two ``layout_q`` / ``layout_k`` parameters are merged into a
249 single ``layout``. The underlying kernel handles all ``cmp_ratio`` values
250 (1 / 4 / 128) directly.
252 The remaining benchmark kwargs (``seqused_q`` / ``seqused_k`` /
253 ``output_idx_offset`` / ``metadata`` / ``max_seqlen_q``) are not yet exposed
254 by this external API and are pinned to ``None`` / ``-1`` inside the DFunction.
256 Args:
257 query: Lightning Indexer query input (Q_index). Must be contiguous.
258 layout='BSND': shape ``(B, S1, N1, D)``; layout='TND': ``(T1, N1, D)``.
259 dtype bfloat16/float16.
260 key: Lightning Indexer key input (K_index). Must be contiguous.
261 Same dtype as query.
262 weights: Weight coefficient (W). shape ``(B, S1, N1)``. Same dtype as query.
263 sparse_count: Number of top-K key tokens to retain (benchmark ``topk``).
264 cu_seq_lens_q: Cumulative query sequence lengths (int32); None for BSND.
265 cu_seq_lens_k: Cumulative key sequence lengths (int32); None for BSND.
266 cmp_residual_k: Per-batch compression residual (original_k_len % cmp_ratio),
267 int32. Affects the valid compressed-key range when ``cmp_ratio != 1``.
268 block_table: Block table for PageAttention (optional).
269 layout: Data layout — 'BSND' (default) or 'TND'. Used for both Q and K.
270 sparse_mode: Sparse mask mode (benchmark ``mask_mode``); 0 = defaultMask.
271 cmp_ratio: Key compression ratio (1 / 4 / 128).
272 return_value: Whether to also output sparse_values (benchmark ``return_value``).
274 Returns:
275 tuple[Tensor, Tensor]: ``(sparse_indices, sparse_values)``.
276 """
277 return _platform.custom_ops.npu_lightning_indexer(
278 query, key, weights, sparse_count,
279 cu_seq_lens_q, cu_seq_lens_k, cmp_residual_k, block_table,
280 layout, sparse_mode, cmp_ratio, return_value,
281 )
284def npu_sparse_flash_mla(
285 query,
286 *,
287 ori_kv: Optional[Tensor] = None,
288 cmp_kv: Optional[Tensor] = None,
289 cmp_sparse_indices: Optional[Tensor] = None,
290 cu_seq_lens_q: Optional[Tensor] = None,
291 cu_seq_lens_ori_kv: Optional[Tensor] = None,
292 cu_seq_lens_cmp_kv: Optional[Tensor] = None,
293 seqused_q: Optional[Tensor] = None,
294 seqused_ori_kv: Optional[Tensor] = None,
295 seqused_cmp_kv: Optional[Tensor] = None,
296 cmp_residual_kv: Optional[Tensor] = None,
297 sinks: Optional[Tensor] = None,
298 softmax_scale: float = 1.0,
299 cmp_ratio: int = 1,
300 ori_mask_mode: int = 4,
301 cmp_mask_mode: int = 3,
302 ori_win_left: int = 127,
303 ori_win_right: int = 0,
304 layout: str = 'BSND',
305 return_softmax_lse: bool = False,
306):
307 """MLA sparse attention (SparseFlashMla).
309 .. warning::
310 This is an experimental API that subject to change or deletion.
312 Computes: O = softmax(Q @ K̃^T · scale) @ Ṽ where K̃ = Ṽ is derived from
313 ``ori_kv``, ``cmp_kv`` and associated sparse indices.
315 The single ``layout`` argument applies to both Q and KV (the DFunction
316 splits it into ``layout_q`` / ``layout_kv`` when calling the kernel). Only
317 ``query`` is positional.
319 Args:
320 query: Query tensor. shape ``(B, S1, N1, D)``, layout BSND.
321 Must be contiguous. dtype bfloat16/float16.
322 ori_kv: Original KV tensor. shape ``(B, S2, 1, D)``. None when absent (band mode).
323 cmp_kv: Compressed KV tensor. shape ``(B, S_cmp, 1, D)``,
324 where ``S_cmp = ceil(S2 / cmp_ratio)``. None when absent.
325 cmp_sparse_indices: Sparse indices for cmp_kv.
326 shape ``(B, S1, 1, K)``, dtype int32. None when cmp_ratio != 4.
327 cu_seq_lens_q: Cumulative query sequence lengths (int32); required for TND.
328 cu_seq_lens_ori_kv: Cumulative ori_kv sequence lengths (int32); for TND.
329 cu_seq_lens_cmp_kv: Cumulative cmp_kv sequence lengths (int32); for TND.
330 seqused_q: Used query sequence lengths (int32); None when absent. For BSND it
331 marks the valid query rows per batch (truncation participates); for TND the
332 query range is governed by cu_seq_lens_q so this is inert.
333 seqused_ori_kv: Used ori_kv sequence lengths (int32); None when absent.
334 seqused_cmp_kv: Used cmp_kv sequence lengths (int32); None when absent.
335 cmp_residual_kv: Per-batch ori_kv-vs-cmp_ratio residual
336 (``ori_len % cmp_ratio``), int32.
337 sinks: Attention-sink tensor. shape ``(N1,)``, dtype float32.
338 softmax_scale: Softmax scaling factor (benchmark default 1.0).
339 cmp_ratio: KV compression ratio (benchmark default 1).
340 ori_mask_mode: Mask mode for q×ori_kv (benchmark default 4 = band).
341 cmp_mask_mode: Mask mode for q×cmp_kv (benchmark default 3 = rightDownCausal).
342 ori_win_left: Band-mask left window (benchmark default 127).
343 ori_win_right: Band-mask right window (benchmark default 0).
344 layout: Data layout for Q and KV — 'BSND' (default) or 'TND'.
345 return_softmax_lse: Whether to also return softmax LSE.
347 Returns:
348 If ``return_softmax_lse=False``: Tensor ``attention_out``,
349 shape ``(B, S1, N1, D)``.
350 If ``return_softmax_lse=True``:
351 tuple[Tensor, Tensor] ``(attention_out, softmax_lse)``.
352 """
353 result = _platform.custom_ops.npu_sparse_flash_mla(
354 query, ori_kv, cmp_kv,
355 cu_seq_lens_q, cu_seq_lens_ori_kv, cu_seq_lens_cmp_kv,
356 None, cmp_sparse_indices, sinks,
357 softmax_scale, cmp_ratio,
358 ori_mask_mode, cmp_mask_mode,
359 ori_win_left, ori_win_right,
360 layout, layout,
361 cmp_residual_kv, seqused_ori_kv, seqused_cmp_kv, seqused_q,
362 )
363 return result if return_softmax_lse else result[0]
366def npu_sparse_flash_mla_grad(
367 query,
368 dout,
369 attn_out,
370 softmax_lse,
371 *,
372 ori_kv: Optional[Tensor] = None,
373 cmp_kv: Optional[Tensor] = None,
374 ori_sparse_indices: Optional[Tensor] = None,
375 cmp_sparse_indices: Optional[Tensor] = None,
376 cu_seq_lens_q: Optional[Tensor] = None,
377 cu_seq_lens_ori_kv: Optional[Tensor] = None,
378 cu_seq_lens_cmp_kv: Optional[Tensor] = None,
379 seqused_q: Optional[Tensor] = None,
380 seqused_ori_kv: Optional[Tensor] = None,
381 seqused_cmp_kv: Optional[Tensor] = None,
382 cmp_residual_kv: Optional[Tensor] = None,
383 ori_topk_length: Optional[Tensor] = None,
384 cmp_topk_length: Optional[Tensor] = None,
385 sinks: Optional[Tensor] = None,
386 softmax_scale: float = 1.0,
387 cmp_ratio: int = 1,
388 ori_mask_mode: int = 4,
389 cmp_mask_mode: int = 3,
390 ori_win_left: int = 127,
391 ori_win_right: int = 0,
392 layout: str = 'BSND',
393) -> Tuple:
394 """MLA sparse-attention backward (SparseFlashMlaGrad), full 6-output form.
396 .. warning::
397 This is an experimental API that subject to change or deletion.
399 Exposes the raw grad kernel so a network-defined custom backward can obtain
400 the ``softmax_l1_norm`` outputs (the main-attention target distribution
401 ``p``) alongside the input gradients, then feed ``p`` straight into
402 :func:`npu_sparse_lightning_indexer_kl_loss_grad` within the same backward.
403 Call this from inside your own autograd function's ``backward`` — it does not
404 build an autograd graph itself.
406 ``metadata`` is not exposed: the grad kernel asserts it is nullptr and
407 re-derives its own tiling internally.
409 Args:
410 query: Query tensor used in the forward. dtype bfloat16/float16.
411 dout: Gradient of the attention output (``grad_attention_out``).
412 attn_out: Attention output from the forward.
413 softmax_lse: Softmax log-sum-exp from the forward.
414 ori_kv: Original KV tensor; None when absent (band mode).
415 cmp_kv: Compressed KV tensor; None when absent.
416 ori_sparse_indices: Sparse indices for ori_kv; None = band mode. Its
417 shape drives ``ori_softmax_l1_norm``.
418 cmp_sparse_indices: Sparse indices for cmp_kv (int32); None when
419 cmp_ratio != 4. Its shape drives ``cmp_softmax_l1_norm``.
420 cu_seq_lens_q: Cumulative query seq lengths (int32); None for BSND.
421 cu_seq_lens_ori_kv: Cumulative ori_kv seq lengths (int32); for TND.
422 cu_seq_lens_cmp_kv: Cumulative cmp_kv seq lengths (int32); for TND.
423 seqused_q: Used query seq lengths (int32); None when absent.
424 seqused_ori_kv: Used ori_kv seq lengths (int32); None when absent.
425 seqused_cmp_kv: Used cmp_kv seq lengths (int32); None when absent.
426 cmp_residual_kv: Per-batch ori_kv-vs-cmp_ratio residual (int32);
427 required for CFA/SCFA (cmp_ratio != 1) with cmp_mask_mode=3.
428 ori_topk_length: Optional ori top-k length; None when absent.
429 cmp_topk_length: Optional cmp top-k length; None when absent.
430 sinks: Attention-sink tensor (float32); None when absent.
431 softmax_scale: Softmax scaling factor (must match the forward).
432 cmp_ratio: KV compression ratio (must match the forward).
433 ori_mask_mode: Mask mode for q×ori_kv (default 4 = band).
434 cmp_mask_mode: Mask mode for q×cmp_kv (default 3 = rightDownCausal).
435 ori_win_left: Band-mask left window (default 127).
436 ori_win_right: Band-mask right window (default 0).
437 layout: Data layout for Q and KV — 'BSND' (default) or 'TND'.
439 Returns:
440 tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
441 ``(d_query, d_ori_kv, d_cmp_kv, d_sinks, ori_softmax_l1_norm,
442 cmp_softmax_l1_norm)``. ``ori/cmp_softmax_l1_norm`` are float32 and
443 share the shape of ``ori/cmp_sparse_indices``; each is the
444 ``reduceG(softmax)/G`` main-attention distribution over the selected
445 tokens. Absent branches yield empty tensors.
446 """
447 return _platform.custom_ops.npu_sparse_flash_mla_grad(
448 query, dout, attn_out, softmax_lse,
449 ori_kv, cmp_kv, ori_sparse_indices, cmp_sparse_indices,
450 cu_seq_lens_q, cu_seq_lens_ori_kv, cu_seq_lens_cmp_kv,
451 seqused_q, seqused_ori_kv, seqused_cmp_kv,
452 cmp_residual_kv, ori_topk_length, cmp_topk_length,
453 sinks, None, # metadata=None → grad kernel self-derives its own tiling
454 softmax_scale, cmp_ratio, ori_mask_mode, cmp_mask_mode,
455 ori_win_left, ori_win_right, layout, layout,
456 )
459def npu_sparse_lightning_indexer_kl_loss_grad(
460 query,
461 key,
462 weights,
463 sparse_indices,
464 attn_softmax_l1_norm,
465 *,
466 cu_seq_lens_q: Optional[Tensor] = None,
467 cu_seq_lens_k: Optional[Tensor] = None,
468 seqused_q: Optional[Tensor] = None,
469 seqused_k: Optional[Tensor] = None,
470 cmp_residual_k: Optional[Tensor] = None,
471 layout: str = 'BSND',
472 mask_mode: int = 3,
473 cmp_ratio: int = 1,
474) -> Tuple:
475 """Compute backward gradients for sparse Lightning Indexer KL loss.
477 .. warning::
478 This is an experimental API that subject to change or deletion.
480 The main-attention target distribution is supplied directly via
481 ``attn_softmax_l1_norm`` (e.g. the ``softmax_l1_norm`` output of
482 :func:`npu_sparse_flash_mla_grad`); the kernel neither recomputes the main
483 attention nor outputs a loss, and returns the indexer-branch softmax as
484 ``softmax_out``.
486 Args:
487 query: Lightning Indexer query input (q̃). Must be contiguous.
488 layout='BSND': shape ``(B, S1, N_qi, D_qi)``;
489 layout='TND': shape ``(T1, N_qi, D_qi)``. dtype bfloat16/float16.
490 key: Lightning Indexer key input (k̃). Must be contiguous.
491 layout='BSND': shape ``(B, S2, N_ki, D_ki)``;
492 layout='TND': shape ``(T2, N_ki, D_ki)``. dtype bfloat16/float16.
493 weights: Weight coefficient (W). Same dtype as query.
494 sparse_indices: Sorted token indices. shape ``(B, S1, 1, K)``, dtype int32.
495 attn_softmax_l1_norm: Main-attention target distribution p (float32),
496 pre-computed by the main-attention branch (e.g. the
497 ``softmax_l1_norm`` output of ``npu_sparse_flash_mla`` backward).
498 cu_seq_lens_q: Cumulative query sequence lengths. shape ``(B+1,)``,
499 dtype int32; None for BSND layout.
500 cu_seq_lens_k: Cumulative key sequence lengths. shape ``(B+1,)``,
501 dtype int32; None for BSND layout.
502 seqused_q: Used query sequence lengths; None when absent.
503 seqused_k: Used key sequence lengths; None when absent.
504 cmp_residual_k: Optional compressed-KV residual.
505 layout: Data layout format — 'TND' (default) or 'BSND'.
506 mask_mode: Sparse mask mode (only 3 supported).
507 cmp_ratio: KV compression ratio.
509 Returns:
510 tuple[Tensor, Tensor, Tensor, Tensor]:
511 ``(d_query, d_key, d_weights, softmax_out)``.
512 """
513 return _platform.custom_ops.npu_sparse_lightning_indexer_kl_loss_grad(
514 query, key, weights, sparse_indices, attn_softmax_l1_norm,
515 cu_seq_lens_q, cu_seq_lens_k, seqused_q, seqused_k, cmp_residual_k,
516 layout, mask_mode, cmp_ratio,
517 )