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.platform import get_platform 

38 

39platform = get_platform() 

40Module = platform.Module 

41 

42 

43class _AsyncSequenceReplicateSlot: 

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

45 

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

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

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

49 compatible with the synchronous DSA module shape. 

50 """ 

51 

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

53 self.device_mesh = device_mesh 

54 self.seq_dim = seq_dim 

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

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

57 self._slots = {} 

58 self._bwd_slots = {} 

59 

60 @staticmethod 

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

62 if _is_tensor_or_dtensor(output): 

63 return output 

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

65 return output[0] 

66 return None 

67 

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

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

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

71 

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

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

74 slot = _make_async_cp_slot( 

75 work, 

76 out_perm, 

77 value, 

78 self.device_mesh, 

79 (Replicate(),), 

80 self.seq_dim, 

81 ) 

82 slot["local"] = local 

83 return slot 

84 

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

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

87 

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

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

90 if module is None: 

91 return 

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

93 

94 def _post_hook(hook_module, hook_inputs, output): 

95 del hook_module, hook_inputs 

96 tensor = self._extract_tensor_output(output) 

97 if tensor is not None: 

98 self.launch(slot_name, tensor) 

99 return output 

100 

101 def _backward_pre_hook(hook_module, grad_output): 

102 del hook_module 

103 return self._producer_bwd_pre_hook(grad_output, bwd_slot) 

104 

105 module.register_forward_hook(_post_hook) 

106 platform.register_full_backward_pre_hook(module, _backward_pre_hook) 

107 

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

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

110 if not bwd_slot: 

111 return grad_output 

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

113 work.wait() 

114 d_local = _allgather_reconstruct(out_perm, gather_dim) 

115 if isinstance(grad_output, tuple): 

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

117 return (d_local,) 

118 

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

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

121 if not _is_tensor_or_dtensor(value): 

122 return 

123 local = self._local_tensor(value) 

124 if not platform.is_tensor(local): 

125 return 

126 if self.world_size <= 1: 

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

128 return 

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

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

131 

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

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

134 if not _is_tensor_or_dtensor(value): 

135 return value 

136 slot = self._slots.get(slot_name) 

137 if slot: 

138 item = slot.pop(0) 

139 if isinstance(item, dict): 

140 local = item["local"] 

141 work, out_perm = _slot_comm(item) 

142 else: 

143 local, work, out_perm = item 

144 if work is None: 

145 return self._wrap_gathered(local, item) 

146 graph_local = self._local_tensor(value) 

147 if not platform.is_tensor(graph_local): 

148 graph_local = local 

149 gathered = platform.differentiable_async_allgather_wait( 

150 graph_local, 

151 work, 

152 out_perm, 

153 self.group, 

154 self.world_size, 

155 self.seq_dim, 

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

157 ) 

158 return self._wrap_gathered(gathered, item) 

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

160 

161 

162class AsyncDSAIndexerContextParallel(DSAIndexerContextParallel): 

163 """Async-first DSA indexer CP. 

164 

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

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

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

168 """ 

169 

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

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

172 cp_mesh = _ensure_1d(device_mesh) 

173 async_state = _AsyncSequenceReplicateSlot(cp_mesh, self.seq_dim) 

174 async_state.register_launch_hook(key_handoff, "key") 

175 specs = self._build_specs( 

176 cp_mesh, 

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

178 ) 

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

180 

181 

182class AsyncDSASparseAttentionContextParallel(DSASparseAttentionContextParallel): 

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

184 

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

186 self, 

187 module: Module, 

188 device_mesh: DeviceMesh, 

189 *, 

190 key_handoff: Optional[Module] = None, 

191 value_handoff: Optional[Module] = None, 

192 key_rope_handoff: Optional[Module] = None, 

193 ) -> Module: 

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

195 cp_mesh = _ensure_1d(device_mesh) 

196 async_state = _AsyncSequenceReplicateSlot(cp_mesh, self.seq_dim) 

197 async_state.register_launch_hook(key_handoff, "key") 

198 async_state.register_launch_hook(value_handoff, "value") 

199 async_state.register_launch_hook(key_rope_handoff, "key_rope") 

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

201 

202 

203class AsyncDSAIndexerLossContextParallel(DSAIndexerLossContextParallel): 

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

205 

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

207 self, 

208 module: Module, 

209 device_mesh: DeviceMesh, 

210 *, 

211 key_handoff: Optional[Module] = None, 

212 key_indexer_handoff: Optional[Module] = None, 

213 key_rope_handoff: Optional[Module] = None, 

214 ) -> Module: 

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

216 cp_mesh = _ensure_1d(device_mesh) 

217 async_state = _AsyncSequenceReplicateSlot(cp_mesh, self.seq_dim) 

218 async_state.register_launch_hook(key_handoff, "key") 

219 async_state.register_launch_hook(key_indexer_handoff, "key_indexer") 

220 async_state.register_launch_hook(key_rope_handoff, "key_rope") 

221 specs = self._build_loss_specs( 

222 cp_mesh, 

223 replicate_fn_map={ 

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

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

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

227 }, 

228 ) 

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