Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / context_parallel / async_dsa_context_parallel.py: 87%

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"""Async context parallel styles for DeepSeek Sparse Attention (DSA).""" 

16from typing import Any, Optional 

17 

18from hyper_parallel.core.context_parallel.async_context_parallel import ( 

19 _allgather_reconstruct, 

20 _launch_async_allgather_seq, 

21 _make_async_cp_slot, 

22 _slot_comm, 

23 _wrap_async_cp_result, 

24) 

25from hyper_parallel.core.context_parallel.context_parallel import _ensure_1d, _localize_foreign_dtensor 

26from hyper_parallel.core.context_parallel.dsa_context_parallel import ( 

27 DSAIndexerContextParallel, 

28 DSAIndexerLossContextParallel, 

29 DSASparseAttentionContextParallel, 

30 _apply_sparse_attention_boundary, 

31 _is_tensor_or_dtensor, 

32 _to_sequence_replicate, 

33) 

34from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

35from hyper_parallel.core.dtensor.dtensor import DTensor 

36from hyper_parallel.core.dtensor.placement_types import Replicate 

37from hyper_parallel.core.context_parallel import utils 

38from hyper_parallel.core.context_parallel.utils import Module 

39 

40 

41 

42class _AsyncSequenceReplicateSlot: 

43 """Per-boundary async sequence-replicate state. 

44 

45 Producer hooks append pre-launched all-gather handles here. Consumer 

46 pre-hooks pop and wait on them. If no producer hook ran, the consumer falls 

47 back to the synchronous sequence-gather path so async DSA styles remain 

48 compatible with the synchronous DSA module shape. 

49 """ 

50 

51 def __init__(self, device_mesh: DeviceMesh, seq_dim: int) -> None: 

52 self.device_mesh = device_mesh 

53 self.seq_dim = seq_dim 

54 self.world_size = device_mesh.mesh.numel() 

55 self.group = device_mesh.get_group() if self.world_size > 1 else None 

56 self._slots = {} 

57 self._bwd_slots = {} 

58 

59 @staticmethod 

60 def _extract_tensor_output(output: Any) -> Any: 

61 if _is_tensor_or_dtensor(output): 

62 return output 

63 if isinstance(output, (tuple, list)) and len(output) == 1 and _is_tensor_or_dtensor(output[0]): 

64 return output[0] 

65 return None 

66 

67 def _local_tensor(self, value: Any) -> Any: 

68 value = _localize_foreign_dtensor(value, self.device_mesh, self.seq_dim) 

69 return value.to_local() if isinstance(value, DTensor) else value 

70 

71 def _make_slot(self, value: Any, local: Any, work, out_perm) -> dict: 

72 """Create an async sequence-replicate slot with layout metadata.""" 

73 slot = _make_async_cp_slot( 

74 work, 

75 out_perm, 

76 value, 

77 self.device_mesh, 

78 (Replicate(),), 

79 self.seq_dim, 

80 ) 

81 slot["local"] = local 

82 return slot 

83 

84 def _wrap_gathered(self, gathered: Any, slot) -> Any: 

85 return _wrap_async_cp_result(gathered, slot, self.device_mesh, (Replicate(),)) 

86 

87 def register_launch_hook(self, module: Optional[Module], slot_name: str) -> None: 

88 """Register a producer-side launch hook when a handoff module exists.""" 

89 if module is None: 

90 return 

91 bwd_slot = self._bwd_slots.setdefault(slot_name, []) 

92 

93 def _post_hook(hook_module, hook_inputs, output): 

94 del hook_module, hook_inputs 

95 tensor = self._extract_tensor_output(output) 

96 if tensor is not None: 

97 self.launch(slot_name, tensor) 

98 return output 

99 

100 def _backward_pre_hook(hook_module, grad_output): 

101 del hook_module 

102 return self._producer_bwd_pre_hook(grad_output, bwd_slot) 

103 

104 module.register_forward_hook(_post_hook) 

105 utils.register_full_backward_pre_hook(module, _backward_pre_hook) 

106 

107 def _producer_bwd_pre_hook(self, grad_output: Any, bwd_slot: list) -> Any: 

108 """Wait deferred reduce-scatter before gradients cross the producer boundary.""" 

109 if not bwd_slot: 

110 return grad_output 

111 work, out_perm, gather_dim = bwd_slot.pop() 

112 work.wait() 

113 d_local = _allgather_reconstruct(out_perm, gather_dim) 

114 if isinstance(grad_output, tuple): 

115 return (d_local,) + grad_output[1:] 

116 return (d_local,) 

117 

118 def launch(self, slot_name: str, value: Any) -> None: 

119 """Launch all-gather for ``value`` and enqueue its handle.""" 

120 if not _is_tensor_or_dtensor(value): 

121 return 

122 local = self._local_tensor(value) 

123 if not utils.is_tensor(local): 

124 return 

125 if self.world_size <= 1: 

126 self._slots.setdefault(slot_name, []).append(self._make_slot(value, local, None, None)) 

127 return 

128 work, out_perm = _launch_async_allgather_seq(local, self.group, self.world_size, self.seq_dim) 

129 self._slots.setdefault(slot_name, []).append(self._make_slot(value, local, work, out_perm)) 

130 

131 def wait(self, slot_name: str, value: Any) -> Any: 

132 """Wait on a pre-launched gather, or fall back to consumer-local launch.""" 

133 if not _is_tensor_or_dtensor(value): 

134 return value 

135 slot = self._slots.get(slot_name) 

136 if slot: 

137 item = slot.pop(0) 

138 if isinstance(item, dict): 

139 local = item["local"] 

140 work, out_perm = _slot_comm(item) 

141 else: 

142 local, work, out_perm = item 

143 if work is None: 

144 return self._wrap_gathered(local, item) 

145 graph_local = self._local_tensor(value) 

146 if not utils.is_tensor(graph_local): 

147 graph_local = local 

148 gathered = utils.differentiable_async_allgather_wait( 

149 graph_local, 

150 work, 

151 out_perm, 

152 self.group, 

153 self.world_size, 

154 self.seq_dim, 

155 self._bwd_slots.setdefault(slot_name, []), 

156 ) 

157 return self._wrap_gathered(gathered, item) 

158 return _to_sequence_replicate(value, self.device_mesh, self.seq_dim) 

159 

160 

161class AsyncDSAIndexerContextParallel(DSAIndexerContextParallel): 

162 """Async-first DSA indexer CP. 

163 

164 When ``key_handoff`` is provided, its forward hook launches key-side 

165 all-gather and the indexer boundary pre-hook waits on it. Without a 

166 handoff, this style falls back to the synchronous key-side gather. 

167 """ 

168 

169 def apply(self, module: Module, device_mesh: DeviceMesh, *, key_handoff: Optional[Module] = None) -> Module: 

170 """Register async DSA indexer CP hooks on ``module`` and optional producer.""" 

171 cp_mesh = _ensure_1d(device_mesh) 

172 async_state = _AsyncSequenceReplicateSlot(cp_mesh, self.seq_dim) 

173 async_state.register_launch_hook(key_handoff, "key") 

174 specs = self._build_specs( 

175 cp_mesh, 

176 key_fn=lambda value: async_state.wait("key", value), 

177 ) 

178 return self._apply_with_specs(module, specs, cp_mesh) 

179 

180 

181class AsyncDSASparseAttentionContextParallel(DSASparseAttentionContextParallel): 

182 """Async-first DSA sparse-attention CP.""" 

183 

184 def apply( # pylint: disable=arguments-differ 

185 self, 

186 module: Module, 

187 device_mesh: DeviceMesh, 

188 *, 

189 key_handoff: Optional[Module] = None, 

190 value_handoff: Optional[Module] = None, 

191 key_rope_handoff: Optional[Module] = None, 

192 ) -> Module: 

193 """Register async DSA sparse-attention CP hooks on ``module`` and producers.""" 

194 cp_mesh = _ensure_1d(device_mesh) 

195 async_state = _AsyncSequenceReplicateSlot(cp_mesh, self.seq_dim) 

196 async_state.register_launch_hook(key_handoff, "key") 

197 async_state.register_launch_hook(value_handoff, "value") 

198 async_state.register_launch_hook(key_rope_handoff, "key_rope") 

199 return _apply_sparse_attention_boundary(self, module, device_mesh, async_state=async_state) 

200 

201 

202class AsyncDSAIndexerLossContextParallel(DSAIndexerLossContextParallel): 

203 """Async-first DSA indexer-loss CP.""" 

204 

205 def apply( # pylint: disable=arguments-differ,too-many-arguments 

206 self, 

207 module: Module, 

208 device_mesh: DeviceMesh, 

209 *, 

210 key_handoff: Optional[Module] = None, 

211 key_indexer_handoff: Optional[Module] = None, 

212 key_rope_handoff: Optional[Module] = None, 

213 ) -> Module: 

214 """Register async DSA indexer-loss CP hooks on ``module`` and producers.""" 

215 cp_mesh = _ensure_1d(device_mesh) 

216 async_state = _AsyncSequenceReplicateSlot(cp_mesh, self.seq_dim) 

217 async_state.register_launch_hook(key_handoff, "key") 

218 async_state.register_launch_hook(key_indexer_handoff, "key_indexer") 

219 async_state.register_launch_hook(key_rope_handoff, "key_rope") 

220 specs = self._build_loss_specs( 

221 cp_mesh, 

222 replicate_fn_map={ 

223 "key": lambda value: async_state.wait("key", value), 

224 "key_indexer": lambda value: async_state.wait("key_indexer", value), 

225 "key_rope": lambda value: async_state.wait("key_rope", value), 

226 }, 

227 ) 

228 return self._apply_with_loss_specs(module, specs, self._get_local_idx(cp_mesh), cp_mesh)