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

1# Copyright 2025-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 

16"""context_parallel: CP collectives, attention adaptations and wrappers. 

17 

18Public surface: the built-in ``@inner_wrapper`` CP schemes 

19(``INNER_WRAPPER_REGISTRY`` + the wrapper factories) and the CP collective 

20primitives (``flex_cp_allgather`` and friends). Split out of 

21components/distributed/cp_utils.py + cp_wrappers.py in stage 4e. 

22""" 

23 

24from hyper_parallel.distributed.context_parallel.collectives import ( 

25 AsyncCPCollective, 

26 async_cp_allgather_launch, 

27 async_ulysses_seq_to_head_launch, 

28 flex_cp_allgather, 

29 hybrid_cp_attention, 

30 ulysses_head_to_seq, 

31 ulysses_seq_to_head, 

32) 

33from hyper_parallel.distributed.context_parallel.wrappers import ( 

34 INNER_WRAPPER_REGISTRY, 

35 INNER_WRAPPER_REQUIREMENTS, 

36 flex_hf_cp_wrapper, 

37 flex_hf_hybrid_cp_wrapper, 

38 flex_hf_ulysses_cp_wrapper, 

39 flex_qkv_cp_wrapper, 

40 flex_qkv_hybrid_cp_wrapper, 

41 flex_qkv_ulysses_cp_wrapper, 

42 mla_dsa_ulysses_cp_wrapper, 

43 sdpa_hf_cp_wrapper, 

44 sdpa_hf_hybrid_cp_wrapper, 

45 sdpa_hf_load_balance_cp_wrapper, 

46 sdpa_hf_ulysses_cp_wrapper, 

47 sdpa_qkv_cp_wrapper, 

48 sdpa_qkv_hybrid_cp_wrapper, 

49 sdpa_qkv_load_balance_cp_wrapper, 

50 sdpa_qkv_ulysses_cp_wrapper, 

51) 

52 

53__all__ = [ 

54 "AsyncCPCollective", 

55 "INNER_WRAPPER_REGISTRY", 

56 "INNER_WRAPPER_REQUIREMENTS", 

57 "async_cp_allgather_launch", 

58 "async_ulysses_seq_to_head_launch", 

59 "flex_cp_allgather", 

60 "flex_hf_cp_wrapper", 

61 "flex_hf_hybrid_cp_wrapper", 

62 "flex_hf_ulysses_cp_wrapper", 

63 "flex_qkv_cp_wrapper", 

64 "flex_qkv_hybrid_cp_wrapper", 

65 "flex_qkv_ulysses_cp_wrapper", 

66 "hybrid_cp_attention", 

67 "mla_dsa_ulysses_cp_wrapper", 

68 "sdpa_hf_cp_wrapper", 

69 "sdpa_hf_hybrid_cp_wrapper", 

70 "sdpa_hf_load_balance_cp_wrapper", 

71 "sdpa_hf_ulysses_cp_wrapper", 

72 "sdpa_qkv_cp_wrapper", 

73 "sdpa_qkv_hybrid_cp_wrapper", 

74 "sdpa_qkv_load_balance_cp_wrapper", 

75 "sdpa_qkv_ulysses_cp_wrapper", 

76 "ulysses_head_to_seq", 

77 "ulysses_seq_to_head", 

78]