Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / custom_ops / experimental / __init__.py: 100%

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 operators for HyperParallel. 

16 

17.. warning:: 

18 This is an experimental API that subject to change or deletion. 

19 

20These operators delegate to the platform-specific ``custom_ops`` interface. 

21On MindSpore they wrap Ascend NPU custom C++ kernels through the 

22``DFunction`` distributed dispatch framework. On PyTorch they raise 

23``NotImplementedError``. 

24 

25Usage:: 

26 

27 from hyper_parallel.custom_ops.experimental import npu_dense_lightning_indexer_softmax_lse 

28 softmax_max, softmax_sum = npu_dense_lightning_indexer_softmax_lse( 

29 query_index, key_index, weights, layout='BSND') 

30 

31When inputs are ``DTensor`` objects, the call is automatically routed 

32through the registered ``DistributedOp`` for layout inference and 

33re-distribution. 

34""" 

35__all__ = [ 

36 "npu_dense_lightning_indexer_grad_kl_loss", 

37 "npu_dense_lightning_indexer_softmax_lse", 

38 "npu_lightning_indexer", 

39 "npu_mhc_post", 

40 "npu_mhc_pre_clamp_sinkhorn", 

41 "npu_mhc_pre_sinkhorn", 

42 "npu_sparse_flash_mla", 

43 "npu_sparse_flash_mla_grad", 

44 "npu_sparse_lightning_indexer_grad_kl_loss", 

45 "npu_sparse_lightning_indexer_kl_loss_grad", 

46] 

47 

48from hyper_parallel.custom_ops.experimental.experimental_ops import ( 

49 npu_dense_lightning_indexer_grad_kl_loss, 

50 npu_dense_lightning_indexer_softmax_lse, 

51 npu_lightning_indexer, 

52 npu_mhc_post, 

53 npu_mhc_pre_clamp_sinkhorn, 

54 npu_mhc_pre_sinkhorn, 

55 npu_sparse_flash_mla, 

56 npu_sparse_flash_mla_grad, 

57 npu_sparse_lightning_indexer_grad_kl_loss, 

58 npu_sparse_lightning_indexer_kl_loss_grad, 

59)