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

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"""context_parallel.wrappers: built-in CP-aware inner attention wrappers (public, explicit injection). 

16 

17Since the explicit-injection rework the framework NEVER picks a CP wrapper 

18automatically: the built-ins below are referenced explicitly — by 

19registry name in an internal plan, by callable, or by a public forward 

20replacement pointing at one of these functions: 

21 

22.. code-block:: yaml 

23 

24 plan_overrides: 

25 - match: "*.self_attn.forward" 

26 when: cp 

27 replace: 

28 kind: method 

29 implementation: hyper_parallel.distributed.context_parallel.wrappers.sdpa_hf_cp_wrapper 

30 

31Wrapper contract (also the Target factory contract — at apply time the 

32Target is built with its declared context, filled by name):: 

33 

34 @inner_wrapper 

35 fn(target_module, mesh, tp_mesh, cp_mesh, ep_mesh, <configured params...>) 

36 

37The ``@inner_wrapper`` decorator is MANDATORY (injection discipline, 

38see recipe_spec.py): the mesh family (``mesh``/``tp_mesh``/``cp_mesh``/ 

39``ep_mesh``) plus the anchor ``target_module`` are REQUIRED context 

40params, ALL filled by the framework at apply time (None for inactive 

41axes) — the user just uses them. 

42Undecorated wrappers fail fast in the resolution chain; *args/**kwargs 

43are forbidden. In-repo wrappers RETURN the replacement — a forward 

44callable or _ForwardRewriteRequest(s) — and the forward rewriter 

45validates, installs and rolls back (05 §15.2.3: module.forward is 

46assigned only by distributed/_builder/forward_rewriter.py); external 

47wrappers may still replace ``target_module.forward`` in place and return 

48None. The replaced forward must accept the original forward's params 

49(validated at apply time). 

50 

51- ``sdpa_qkv_cp_wrapper``: NeMo convention ``forward(q, k, v, ...)`` — explicit 

52 K/V all-gather + D-04 offset-aware causal mask (fixes the is_causal 

53 alignment error under CP); 

54- ``sdpa_hf_cp_wrapper``: HF convention ``forward(hidden_states, ...)`` — 

55 primitive interception of ``F.scaled_dot_product_attention`` (reuses HF 

56 projections/RoPE), with misfire detection (raises if no call is 

57 intercepted); 

58- ``flex_qkv_cp_wrapper`` / ``flex_hf_cp_wrapper``: the two isomorphic 

59 FlexAttention entries (block_mask must be built for the global kv length). 

60- ``*_ulysses_cp_wrapper``: Pure Ulysses variants that exchange Q/K/V from 

61 sequence shards to head shards before attention and restore the local 

62 sequence shard afterward. 

63- ``*_hybrid_cp_wrapper``: local-tensor Hybrid variants that run Ulysses 

64 all-to-all inside each subgroup and K/V all-gather across the complementary 

65 Colossal subgroup. 

66- ``sdpa_*_load_balance_cp_wrapper``: local-tensor Colossal Head-Tail variants 

67 for Q exchange, K/V all-gather, dual SDPA execution, output restoration, 

68 and backward communication. 

69- ``gdn_ulysses_cp_wrapper``: primitive interception for Gated DeltaNet, 

70 with a causal-Conv1d halo and Ulysses exchange around the 

71 Gated Delta Rule interface. 

72 

73Users may register their own named schemes:: 

74 

75 INNER_WRAPPER_REGISTRY["my_flash"] = my_fn 

76 

77after which ``spec.inner_wrapper="my_flash"`` references it by name. 

78 

79The ``is_*_attention`` helpers are public utilities for authors of custom 

80wrappers (choosing a scheme programmatically); the framework itself no 

81longer dispatches heuristically on them. 

82""" 

83 

84import functools 

85import importlib 

86import inspect 

87import logging 

88from typing import Any, Callable 

89 

90import torch # pylint: disable=forbidden-backend-import 

91import torch.nn.functional as F 

92from torch.nn import Module 

93 

94from hyper_parallel.distributed._builder.forward_rewriter import ( 

95 _ForwardRewriteRequest, 

96) 

97from hyper_parallel.distributed.context_parallel.collectives import ( 

98 _build_hybrid_cp_submeshes, 

99 _ULYSSES_WRAPPED_FLAG, 

100 _UlyssesContext, 

101 _slice_sequence, 

102 _validate_ulysses_head_counts, 

103 flex_cp_allgather, 

104 hybrid_cp_attention, 

105 ulysses_head_to_seq, 

106 ulysses_seq_to_head, 

107) 

108from hyper_parallel.distributed.context_parallel.attention import ( 

109 _cp_offset_causal_mask, 

110 _dsa_cp_alltoall, 

111 _mla_cp_alltoall, 

112 _mome_cp_halo_exchange, 

113 head_tail_load_balance_attention, 

114) 

115from hyper_parallel.distributed.context_parallel.gated_delta_net import install_gdn_ulysses_cp 

116from hyper_parallel.distributed.recipe_spec import ( 

117 inner_wrapper, 

118) 

119 

120logger = logging.getLogger(__name__) 

121 

122 

123# ──────────────────────────────────────────────────────────────────────────── 

124# Style-detection helpers (public utilities for custom wrapper authors) 

125# ──────────────────────────────────────────────────────────────────────────── 

126 

127def _attn_implementation(module): 

128 cfg = getattr(module, "config", None) 

129 impl = getattr(cfg, "_attn_implementation", None) 

130 if impl is None and isinstance(cfg, dict): 

131 impl = cfg.get("attn_implementation") 

132 return impl 

133 

134 

135def is_sdpa_attention(module: Any) -> bool: 

136 """True when the module is configured for or named as SDPA attention.""" 

137 impl = _attn_implementation(module) 

138 return (impl == "sdpa") or ("SdpaAttention" in type(module).__name__) 

139 

140 

141def is_flex_attention(module: Any) -> bool: 

142 """True when the module is configured for or named as FlexAttention.""" 

143 impl = _attn_implementation(module) 

144 return (impl == "flex_attention") or ("FlexAttention" in type(module).__name__) 

145 

146 

147def is_hf_style_attention(module: Any) -> bool: 

148 """HF style (forward(hidden_states,...), projections inside forward) -> the primitive-interception path.""" 

149 has_proj = (hasattr(module, "q_proj") and hasattr(module, "k_proj") 

150 and hasattr(module, "v_proj")) 

151 if not has_proj: 

152 return False 

153 try: 

154 sig = inspect.signature(module.forward) 

155 first_param = next(iter(sig.parameters.values()), None) 

156 return first_param is not None and first_param.name == "hidden_states" 

157 except (ValueError, TypeError): 

158 return type(module).__name__.endswith("Attention") 

159 

160 

161# ──────────────────────────────────────────────────────────────────────────── 

162# CP-aware SDPA core (K/V all-gather + D-04 offset causal mask) 

163# ──────────────────────────────────────────────────────────────────────────── 

164 

165def _prepare_cp_sdpa_kwargs( 

166 kwargs: dict[str, Any], *, q_len: int, kv_len: int, 

167 query_offset: int, device: Any) -> dict[str, Any]: 

168 """Adapt a complete SDPA mask to one CP-local query interval.""" 

169 call_kwargs = dict(kwargs) 

170 attention_mask = call_kwargs.get("attn_mask") 

171 if attention_mask is not None: 

172 if not isinstance(attention_mask, torch.Tensor): 

173 raise TypeError( 

174 "CP SDPA attn_mask must be a Tensor or None, got " 

175 f"{type(attention_mask).__name__}" 

176 ) 

177 if attention_mask.ndim < 2: 

178 raise ValueError( 

179 "CP SDPA attn_mask must expose query and KV dimensions, got " 

180 f"shape {tuple(attention_mask.shape)}" 

181 ) 

182 if attention_mask.shape[-1] != kv_len: 

183 raise ValueError( 

184 "CP SDPA attn_mask must cover the complete KV sequence: " 

185 f"mask KV length={attention_mask.shape[-1]}, expected {kv_len}" 

186 ) 

187 mask_q_len = attention_mask.shape[-2] 

188 if mask_q_len != q_len: 

189 if mask_q_len < query_offset + q_len: 

190 raise ValueError( 

191 "CP SDPA attn_mask does not cover the required query range " 

192 f"[{query_offset}, {query_offset + q_len}); mask query " 

193 f"length={mask_q_len}" 

194 ) 

195 attention_mask = attention_mask.narrow( 

196 -2, query_offset, q_len 

197 ) 

198 call_kwargs["attn_mask"] = attention_mask 

199 if call_kwargs.get("is_causal"): 

200 call_kwargs["is_causal"] = False 

201 return call_kwargs 

202 

203 if call_kwargs.get("is_causal"): 

204 call_kwargs["is_causal"] = False 

205 call_kwargs["attn_mask"] = _cp_offset_causal_mask( 

206 q_len, kv_len, query_offset, device 

207 ) 

208 return call_kwargs 

209 

210 

211def _prepare_head_tail_sdpa_kwargs( 

212 kwargs: dict[str, Any], query: Any, key: Any, 

213 cp_mesh: Any) -> tuple[dict[str, Any], dict[str, Any]]: 

214 """Prepare complete-mask query rows for both Head-Tail attention calls.""" 

215 local_q_len = query.shape[2] 

216 if local_q_len % 2: 

217 raise ValueError( 

218 "Head-Tail load balance requires an even local Q sequence " 

219 f"length, got {local_q_len}" 

220 ) 

221 half_q_len = local_q_len // 2 

222 local_rank = cp_mesh.get_local_rank() 

223 peer_rank = cp_mesh.size() - 1 - local_rank 

224 global_kv_len = key.shape[2] * cp_mesh.size() 

225 keep_kwargs = _prepare_cp_sdpa_kwargs( 

226 kwargs, 

227 q_len=half_q_len, 

228 kv_len=global_kv_len, 

229 query_offset=local_rank * local_q_len, 

230 device=query.device, 

231 ) 

232 peer_kwargs = _prepare_cp_sdpa_kwargs( 

233 kwargs, 

234 q_len=half_q_len, 

235 kv_len=global_kv_len, 

236 query_offset=peer_rank * local_q_len + half_q_len, 

237 device=query.device, 

238 ) 

239 return keep_kwargs, peer_kwargs 

240 

241 

242def _prepare_hybrid_sdpa_kwargs( 

243 kwargs: dict[str, Any], query: Any, key: Any, 

244 cp_mesh: Any, ulysses_degree: int) -> dict[str, Any]: 

245 """Prepare complete-mask query rows for Hybrid attention.""" 

246 _, colossal_mesh = _build_hybrid_cp_submeshes(cp_mesh, ulysses_degree) 

247 hybrid_q_len = query.shape[2] * ulysses_degree 

248 return _prepare_cp_sdpa_kwargs( 

249 kwargs, 

250 q_len=hybrid_q_len, 

251 kv_len=key.shape[2] * cp_mesh.size(), 

252 query_offset=colossal_mesh.get_local_rank() * hybrid_q_len, 

253 device=query.device, 

254 ) 

255 

256 

257def _cp_sdpa_call(orig_sdpa, cp_mesh, q, k, v, kwargs): 

258 """CP-aware SDPA: K/V all-gather + D-04 offset-aware causal mask.""" 

259 cp_dim = 2 # sequence dim of the [B, N, S, H] layout 

260 global_k, global_v = flex_cp_allgather( 

261 k.contiguous(), v.contiguous(), cp_dim, cp_mesh) 

262 if cp_mesh.size() > 1: 

263 kwargs = _prepare_cp_sdpa_kwargs( 

264 kwargs, 

265 q_len=q.shape[cp_dim], 

266 kv_len=global_k.shape[cp_dim], 

267 query_offset=cp_mesh.get_local_rank() * q.shape[cp_dim], 

268 device=q.device, 

269 ) 

270 return orig_sdpa(q, global_k, global_v, **kwargs) 

271 

272 

273def _ulysses_attention_call( 

274 attention_fn, cp_mesh, query, key, value, kwargs, 

275 *, seq_dim=2, head_dim=1): 

276 """Run attention in the Pure Ulysses full-sequence/head-sharded layout.""" 

277 _validate_ulysses_head_counts( 

278 query, key, value, cp_mesh.size(), "Pure Ulysses CP" 

279 ) 

280 query, key, value = ( 

281 ulysses_seq_to_head(tensor, seq_dim, head_dim, cp_mesh) 

282 for tensor in (query, key, value) 

283 ) 

284 output = attention_fn(query, key, value, **kwargs) 

285 return ulysses_head_to_seq(output, seq_dim, head_dim, cp_mesh) 

286 

287 

288def _ulysses_qkv_forward(original_forward, cp_mesh, args, kwargs): 

289 """Preserve a QKV forward signature while applying Pure Ulysses.""" 

290 signature = inspect.signature(original_forward) 

291 bound = signature.bind(*args, **kwargs) 

292 names = list(signature.parameters) 

293 if len(names) < 3: 

294 raise TypeError("Ulysses QKV wrapper requires at least three forward inputs") 

295 try: 

296 query, key, value = ( 

297 bound.arguments[names[index]] for index in range(3) 

298 ) 

299 except KeyError as exc: 

300 raise TypeError( 

301 "Ulysses QKV wrapper requires query, key, and value inputs" 

302 ) from exc 

303 _validate_ulysses_head_counts( 

304 query, key, value, cp_mesh.size(), "Pure Ulysses CP" 

305 ) 

306 query, key, value = ( 

307 ulysses_seq_to_head(tensor, 2, 1, cp_mesh) 

308 for tensor in (query, key, value) 

309 ) 

310 bound.arguments[names[0]] = query 

311 bound.arguments[names[1]] = key 

312 bound.arguments[names[2]] = value 

313 output = original_forward(*bound.args, **bound.kwargs) 

314 return ulysses_head_to_seq(output, 2, 1, cp_mesh) 

315 

316 

317def _require_ulysses_cp_mesh(cp_mesh, wrapper_name): 

318 """Validate the communication context required by Pure Ulysses.""" 

319 if cp_mesh is None or cp_mesh.size() <= 1: 

320 raise ValueError( 

321 f"Ulysses wrapper {wrapper_name!r} requires an active CP mesh" 

322 ) 

323 

324 

325# ──────────────────────────────────────────────────────────────────────────── 

326# The four built-in wrappers 

327# ──────────────────────────────────────────────────────────────────────────── 

328 

329@inner_wrapper 

330def sdpa_qkv_cp_wrapper( 

331 target_module: Module, mesh: Any, tp_mesh: Any, 

332 cp_mesh: Any, ep_mesh: Any) -> Callable: 

333 """NeMo/Megatron SDPA path (registry "sdpa_qkv"): explicit all-gather K/V. 

334 

335 Assumes the inner_attention.forward(q,k,v,...) signature convention. 

336 **Local-only** (injection discipline): the dual-mode adapter converts 

337 DTensor inputs to local and re-wraps the output — this wrapper never 

338 touches DTensor. Since it wraps an inner submodule, the plan MUST 

339 declare ``inner_out_src: "first_input"`` (its output layout == q's 

340 layout). cp_mesh is framework-filled context; fail fast when the plan 

341 has no cp axis. 

342 """ 

343 

344 del mesh, tp_mesh, ep_mesh 

345 if cp_mesh is None: 

346 raise ValueError( 

347 "The in-repo CP wrapper reference implementation 'sdpa_qkv' requires an " 

348 "active cp axis (the communication domain for K/V all-gather), but the " 

349 "framework-filled cp_mesh is None (the current plan has no cp axis) — " 

350 "without CP, use a custom @inner_wrapper wrapper instead (the " 

351 "cp_mesh=None semantics are then your own responsibility), or use the " 

352 "local_compute_fn channel (see " 

353 "examples/distributed/perf_replacement.py)") 

354 

355 original_forward = target_module.forward 

356 

357 @functools.wraps(original_forward) 

358 def cp_forward(q: Any, k: Any, v: Any, **kwargs: Any) -> Any: 

359 """All-gather K/V along CP, then run the original QKV SDPA forward.""" 

360 return _cp_sdpa_call(original_forward, cp_mesh, q, k, v, kwargs) 

361 

362 return cp_forward 

363 

364 

365@inner_wrapper 

366def flex_qkv_cp_wrapper( 

367 target_module: Module, mesh: Any, tp_mesh: Any, 

368 cp_mesh: Any, ep_mesh: Any) -> Callable: 

369 """NeMo/Megatron FlexAttention path (registry "flex_qkv"): explicit all-gather K/V. 

370 

371 **Local-only** (see sdpa_qkv_cp_wrapper); requires 

372 ``inner_out_src: "first_input"`` in the plan. Constraint: the 

373 block_mask must be built for the global kv length. cp_mesh is 

374 framework-filled context; fail fast when the plan has no cp axis. 

375 """ 

376 

377 del mesh, tp_mesh, ep_mesh 

378 if cp_mesh is None: 

379 raise ValueError( 

380 "The in-repo CP wrapper reference implementation 'flex_qkv' requires an " 

381 "active cp axis (the communication domain for K/V all-gather), but the " 

382 "framework-filled cp_mesh is None (the current plan has no cp axis) — " 

383 "without CP, use a custom @inner_wrapper wrapper instead (the " 

384 "cp_mesh=None semantics are then your own responsibility), or use the " 

385 "local_compute_fn channel (see " 

386 "examples/distributed/perf_replacement.py)") 

387 

388 original_forward = target_module.forward 

389 

390 @functools.wraps(original_forward) 

391 def cp_forward(q: Any, k: Any, v: Any, **kwargs: Any) -> Any: 

392 """All-gather K/V along CP, then run the original QKV Flex forward.""" 

393 global_k, global_v = flex_cp_allgather( 

394 k.contiguous(), v.contiguous(), 2, cp_mesh) 

395 return original_forward(q, global_k, global_v, **kwargs) 

396 

397 return cp_forward 

398 

399 

400@inner_wrapper 

401def sdpa_hf_cp_wrapper( 

402 target_module: Module, mesh: Any, tp_mesh: Any, 

403 cp_mesh: Any, ep_mesh: Any) -> Callable: 

404 """HF standard SDPA path: forward(hidden_states,...) -> primitive interception (05 §4.4.2). 

405 

406 **Local-only** (see sdpa_qkv_cp_wrapper): the adapter handles DTensor 

407 unwrap/param-unwrap/rewrap; the rewrap uses the boundary's declared 

408 out_src (this wrapper targets the boundary module itself). The 

409 primitive interception is a temporary global function replacement 

410 (restored via try/finally) and is not thread-safe; it is safe under 

411 single-process SPMD training (consistent with the TorchTitan CP 

412 implementation). 

413 Misfire detection: the wrapper may be applied to a module that does not 

414 call F.sdpa -- if the primitive did not intercept a single call, the K/V 

415 were not gathered (a silent numerical error), so raise a RuntimeError 

416 immediately. 

417 """ 

418 

419 del mesh, tp_mesh, ep_mesh 

420 if cp_mesh is None: 

421 raise ValueError( 

422 "The in-repo CP wrapper reference implementation 'sdpa_hf' requires an " 

423 "active cp axis (the communication domain for K/V all-gather), but the " 

424 "framework-filled cp_mesh is None (the current plan has no cp axis) — " 

425 "without CP, use a custom @inner_wrapper wrapper instead (the " 

426 "cp_mesh=None semantics are then your own responsibility), or use the " 

427 "local_compute_fn channel (see " 

428 "examples/distributed/perf_replacement.py)") 

429 

430 setattr(target_module, "_hp_sdpa_hf_cp_mesh", cp_mesh) 

431 setattr(target_module, "_hp_sdpa_hf_original_forward", target_module.forward) 

432 

433 @functools.wraps(target_module.forward) 

434 def cp_forward(hidden_states: Any, *args: Any, **kwargs: Any) -> Any: 

435 """Delegate to the source-visible SDPA helper shared with Codegen.""" 

436 return sdpa_hf_cp_forward(target_module, hidden_states, *args, **kwargs) 

437 

438 return cp_forward 

439 

440 

441def sdpa_hf_cp_forward( 

442 target_module: Module, hidden_states: Any, 

443 *args: Any, **kwargs: Any) -> Any: 

444 """Run an HF attention forward while intercepting SDPA for CP.""" 

445 cp_mesh = getattr(target_module, "_hp_sdpa_hf_cp_mesh") 

446 original_forward = getattr(target_module, "_forward_impl", None) 

447 if original_forward is None: 

448 original_forward = getattr(target_module, "_hp_sdpa_hf_original_forward") 

449 orig_sdpa = F.scaled_dot_product_attention 

450 fired = {"hit": False} 

451 

452 def cp_aware_sdpa(q: Any, k: Any, v: Any, **kw: Any) -> Any: 

453 """CP-aware SDPA replacement: all-gather K/V plus the D-04 mask.""" 

454 fired["hit"] = True 

455 return _cp_sdpa_call(orig_sdpa, cp_mesh, q, k, v, kw) 

456 

457 F.scaled_dot_product_attention = cp_aware_sdpa 

458 try: 

459 out = original_forward(hidden_states, *args, **kwargs) 

460 finally: 

461 F.scaled_dot_product_attention = orig_sdpa 

462 if not fired["hit"]: 

463 raise RuntimeError( 

464 f"CP wrapper 'sdpa_hf' did not intercept any " 

465 f"F.scaled_dot_product_attention call on " 

466 f"{type(target_module).__name__} -- the wrapper does " 

467 f"not match the module implementation (K/V were not " 

468 f"all-gathered; continuing would produce silent numerical " 

469 f"errors). Please explicitly set inner_wrapper='sdpa_qkv' " 

470 f"(the (q,k,v) convention), or provide a custom inner_wrapper " 

471 f"callable") 

472 return out 

473 

474 

475@inner_wrapper 

476def flex_hf_cp_wrapper( 

477 target_module: Module, mesh: Any, tp_mesh: Any, 

478 cp_mesh: Any, ep_mesh: Any) -> Callable: 

479 """HF standard FlexAttention path: intercept flex_attention (same structure as the SDPA path). 

480 

481 **Local-only** (see sdpa_qkv_cp_wrapper). Constraint: 

482 score_mod/block_mask pass through verbatim via kwargs -- under CP, 

483 kv_len changes from S/cp to S, so the block_mask must be built for the 

484 **global kv length** (constructed on the full sequence in the data 

485 pipeline / model side), otherwise shapes and semantics are misaligned. 

486 The wrapper does not validate this. 

487 Misfire detection is the same as 'sdpa_hf': if no flex_attention call is 

488 intercepted, raise a RuntimeError. 

489 """ 

490 

491 del mesh, tp_mesh, ep_mesh 

492 if cp_mesh is None: 

493 raise ValueError( 

494 "The in-repo CP wrapper reference implementation 'flex_hf' requires an " 

495 "active cp axis (the communication domain for K/V all-gather), but the " 

496 "framework-filled cp_mesh is None (the current plan has no cp axis) — " 

497 "without CP, use a custom @inner_wrapper wrapper instead (the " 

498 "cp_mesh=None semantics are then your own responsibility), or use the " 

499 "local_compute_fn channel (see " 

500 "examples/distributed/perf_replacement.py)") 

501 

502 original_forward = target_module.forward 

503 # FlexAttention is optional on older supported PyTorch versions. 

504 flex_attention_module = importlib.import_module( 

505 "torch.nn.attention.flex_attention") 

506 original_flex_attention = flex_attention_module.flex_attention 

507 

508 @functools.wraps(original_forward) 

509 def cp_forward(hidden_states: Any, *args: Any, **kwargs: Any) -> Any: 

510 """Run the HF forward with flex_attention temporarily replaced.""" 

511 fired = {"hit": False} 

512 

513 def cp_aware_flex(q: Any, k: Any, v: Any, **kw: Any) -> Any: 

514 """CP-aware FlexAttention replacement: all-gather K/V first.""" 

515 fired["hit"] = True 

516 global_k, global_v = flex_cp_allgather( 

517 k.contiguous(), v.contiguous(), 2, cp_mesh) 

518 return original_flex_attention(q, global_k, global_v, **kw) 

519 

520 flex_attention_module.flex_attention = cp_aware_flex 

521 try: 

522 out = original_forward(hidden_states, *args, **kwargs) 

523 finally: 

524 flex_attention_module.flex_attention = original_flex_attention 

525 if not fired["hit"]: 

526 raise RuntimeError( 

527 f"CP wrapper 'flex_hf' did not intercept any flex_attention " 

528 f"call on {type(target_module).__name__} -- the wrapper " 

529 f"does not match the module implementation (K/V were " 

530 f"not all-gathered; continuing would produce silent numerical " 

531 f"errors). Please explicitly set inner_wrapper='flex_qkv' " 

532 f"(the (q,k,v) convention), or provide a custom inner_wrapper " 

533 f"callable") 

534 return out 

535 

536 return cp_forward 

537 

538 

539# ──────────────────────────────────────────────────────────────────────────── 

540# Colossal Head-Tail load-balance wrappers (explicit plan_overrides methods) 

541# ──────────────────────────────────────────────────────────────────────────── 

542 

543@inner_wrapper 

544def sdpa_qkv_load_balance_cp_wrapper( 

545 target_module: Module, mesh: Any, tp_mesh: Any, 

546 cp_mesh: Any, ep_mesh: Any) -> Callable: 

547 """Apply Colossal Head-Tail load balancing to a QKV-style SDPA module. 

548 

549 Args: 

550 target_module: Attention module whose first three inputs are Q/K/V. 

551 mesh: Framework-owned root mesh; unused by this wrapper. 

552 tp_mesh: Framework-owned TP mesh; unused by this wrapper. 

553 cp_mesh: Framework-owned CP communication mesh. 

554 ep_mesh: Framework-owned EP mesh; unused by this wrapper. 

555 """ 

556 del mesh, tp_mesh, ep_mesh 

557 if cp_mesh is None or cp_mesh.size() <= 1: 

558 raise ValueError( 

559 "sdpa_qkv_load_balance_cp_wrapper requires an active CP mesh" 

560 ) 

561 original_forward = target_module.forward 

562 

563 @functools.wraps(original_forward) 

564 def cp_forward(q: Any, k: Any, v: Any, **kwargs: Any) -> Any: 

565 """Route local Q/K/V tensors through Head-Tail communication.""" 

566 def attention_call( 

567 query: Any, key: Any, value: Any, 

568 call_kwargs: dict[str, Any]) -> Any: 

569 """Invoke the original QKV forward with transformed tensors.""" 

570 return original_forward(query, key, value, **call_kwargs) 

571 

572 keep_kwargs, peer_kwargs = _prepare_head_tail_sdpa_kwargs( 

573 kwargs, q, k, cp_mesh 

574 ) 

575 return head_tail_load_balance_attention( 

576 attention_call, 

577 q, 

578 k, 

579 v, 

580 keep_kwargs, 

581 cp_mesh, 

582 peer_attention_kwargs=peer_kwargs, 

583 ) 

584 

585 return cp_forward 

586 

587 

588@inner_wrapper 

589def sdpa_hf_load_balance_cp_wrapper( 

590 target_module: Module, mesh: Any, tp_mesh: Any, 

591 cp_mesh: Any, ep_mesh: Any) -> Callable: 

592 """Apply Colossal Head-Tail load balancing at an HF SDPA primitive. 

593 

594 Args: 

595 target_module: HF-style attention containing an SDPA primitive call. 

596 mesh: Framework-owned root mesh; unused by this wrapper. 

597 tp_mesh: Framework-owned TP mesh; unused by this wrapper. 

598 cp_mesh: Framework-owned CP communication mesh. 

599 ep_mesh: Framework-owned EP mesh; unused by this wrapper. 

600 """ 

601 del mesh, tp_mesh, ep_mesh 

602 if cp_mesh is None or cp_mesh.size() <= 1: 

603 raise ValueError( 

604 "sdpa_hf_load_balance_cp_wrapper requires an active CP mesh" 

605 ) 

606 original_forward = target_module.forward 

607 original_sdpa = F.scaled_dot_product_attention 

608 

609 @functools.wraps(original_forward) 

610 def cp_forward( 

611 hidden_states: Any, *args: Any, **kwargs: Any) -> Any: 

612 """Intercept the module's SDPA primitive for one forward call.""" 

613 fired = {"hit": False} 

614 

615 def load_balanced_sdpa( 

616 q: Any, k: Any, v: Any, 

617 **call_kwargs: Any) -> Any: 

618 """Run one intercepted SDPA call with Head-Tail communication.""" 

619 fired["hit"] = True 

620 keep_kwargs, peer_kwargs = _prepare_head_tail_sdpa_kwargs( 

621 call_kwargs, q, k, cp_mesh 

622 ) 

623 

624 def _call_original_sdpa(query, key, value, attention_kwargs): 

625 return original_sdpa(query, key, value, **attention_kwargs) 

626 

627 return head_tail_load_balance_attention( 

628 _call_original_sdpa, 

629 q, 

630 k, 

631 v, 

632 keep_kwargs, 

633 cp_mesh, 

634 peer_attention_kwargs=peer_kwargs, 

635 ) 

636 

637 F.scaled_dot_product_attention = load_balanced_sdpa 

638 try: 

639 output = original_forward(hidden_states, *args, **kwargs) 

640 finally: 

641 F.scaled_dot_product_attention = original_sdpa 

642 if not fired["hit"]: 

643 raise RuntimeError( 

644 "CP wrapper 'sdpa_hf_load_balance' did not intercept any " 

645 "F.scaled_dot_product_attention call on " 

646 f"{type(target_module).__name__}; choose the QKV wrapper or " 

647 "provide a matching custom inner_wrapper" 

648 ) 

649 return output 

650 

651 return cp_forward 

652 

653 

654# ──────────────────────────────────────────────────────────────────────────── 

655# Pure Ulysses attention wrappers (explicit plan_overrides methods) 

656# ──────────────────────────────────────────────────────────────────────────── 

657 

658@inner_wrapper 

659def sdpa_qkv_ulysses_cp_wrapper( 

660 target_module: Module, mesh: Any, tp_mesh: Any, 

661 cp_mesh: Any, ep_mesh: Any) -> Callable: 

662 """Apply Pure Ulysses to a separated ``forward(query, key, value, ...)``.""" 

663 del mesh, tp_mesh, ep_mesh 

664 _require_ulysses_cp_mesh(cp_mesh, "sdpa_qkv_ulysses") 

665 original_forward = target_module.forward 

666 

667 @functools.wraps(original_forward) 

668 def cp_forward(*args: Any, **kwargs: Any) -> Any: 

669 """Run the original QKV forward in the Ulysses head-sharded layout.""" 

670 return _ulysses_qkv_forward(original_forward, cp_mesh, args, kwargs) 

671 

672 return cp_forward 

673 

674 

675@inner_wrapper 

676def flex_qkv_ulysses_cp_wrapper( 

677 target_module: Module, mesh: Any, tp_mesh: Any, 

678 cp_mesh: Any, ep_mesh: Any) -> Callable: 

679 """Apply Pure Ulysses to a separated FlexAttention QKV forward.""" 

680 del mesh, tp_mesh, ep_mesh 

681 _require_ulysses_cp_mesh(cp_mesh, "flex_qkv_ulysses") 

682 original_forward = target_module.forward 

683 

684 @functools.wraps(original_forward) 

685 def cp_forward(*args: Any, **kwargs: Any) -> Any: 

686 """Run the original FlexAttention QKV forward in the Ulysses layout.""" 

687 return _ulysses_qkv_forward(original_forward, cp_mesh, args, kwargs) 

688 

689 return cp_forward 

690 

691 

692@inner_wrapper 

693def sdpa_hf_ulysses_cp_wrapper( 

694 target_module: Module, mesh: Any, tp_mesh: Any, 

695 cp_mesh: Any, ep_mesh: Any) -> Callable: 

696 """Apply Pure Ulysses by intercepting HF SDPA primitive calls.""" 

697 del mesh, tp_mesh, ep_mesh 

698 _require_ulysses_cp_mesh(cp_mesh, "sdpa_hf_ulysses") 

699 original_forward = target_module.forward 

700 original_sdpa = F.scaled_dot_product_attention 

701 

702 @functools.wraps(original_forward) 

703 def cp_forward(*args: Any, **kwargs: Any) -> Any: 

704 """Run the HF forward with F.sdpa temporarily Ulysses-exchanged.""" 

705 fired = {"hit": False} 

706 

707 def ulysses_sdpa( 

708 query: Any, key: Any, value: Any, 

709 **attention_kwargs: Any) -> Any: 

710 """Run one intercepted SDPA call in the Ulysses layout.""" 

711 fired["hit"] = True 

712 return _ulysses_attention_call( 

713 original_sdpa, 

714 cp_mesh, 

715 query, 

716 key, 

717 value, 

718 attention_kwargs, 

719 ) 

720 

721 F.scaled_dot_product_attention = ulysses_sdpa 

722 try: 

723 output = original_forward(*args, **kwargs) 

724 finally: 

725 F.scaled_dot_product_attention = original_sdpa 

726 if not fired["hit"]: 

727 raise RuntimeError( 

728 "CP wrapper 'sdpa_hf_ulysses' did not intercept any " 

729 "F.scaled_dot_product_attention call; the selected wrapper " 

730 "does not match the module implementation" 

731 ) 

732 return output 

733 

734 return cp_forward 

735 

736 

737@inner_wrapper 

738def flex_hf_ulysses_cp_wrapper( 

739 target_module: Module, mesh: Any, tp_mesh: Any, 

740 cp_mesh: Any, ep_mesh: Any) -> Callable: 

741 """Apply Pure Ulysses by intercepting HF FlexAttention primitive calls.""" 

742 del mesh, tp_mesh, ep_mesh 

743 _require_ulysses_cp_mesh(cp_mesh, "flex_hf_ulysses") 

744 original_forward = target_module.forward 

745 flex_attention_module = importlib.import_module( 

746 "torch.nn.attention.flex_attention" 

747 ) 

748 original_flex_attention = flex_attention_module.flex_attention 

749 

750 @functools.wraps(original_forward) 

751 def cp_forward(*args: Any, **kwargs: Any) -> Any: 

752 """Run the HF forward with flex_attention temporarily Ulysses-exchanged.""" 

753 fired = {"hit": False} 

754 

755 def ulysses_flex( 

756 query: Any, key: Any, value: Any, 

757 **attention_kwargs: Any) -> Any: 

758 """Run one intercepted FlexAttention call in the Ulysses layout.""" 

759 fired["hit"] = True 

760 return _ulysses_attention_call( 

761 original_flex_attention, 

762 cp_mesh, 

763 query, 

764 key, 

765 value, 

766 attention_kwargs, 

767 ) 

768 

769 flex_attention_module.flex_attention = ulysses_flex 

770 try: 

771 output = original_forward(*args, **kwargs) 

772 finally: 

773 flex_attention_module.flex_attention = original_flex_attention 

774 if not fired["hit"]: 

775 raise RuntimeError( 

776 "CP wrapper 'flex_hf_ulysses' did not intercept any " 

777 "flex_attention call; the selected wrapper does not match " 

778 "the module implementation" 

779 ) 

780 return output 

781 

782 return cp_forward 

783 

784 

785# ──────────────────────────────────────────────────────────────────────────── 

786# Hybrid attention wrappers (explicit Target with wrapper-local degree) 

787# ──────────────────────────────────────────────────────────────────────────── 

788 

789def _validate_hybrid_config( 

790 cp_mesh: Any, ulysses_degree: int, wrapper_name: str) -> None: 

791 """Validate Hybrid topology at wrapper installation time.""" 

792 if cp_mesh is None or cp_mesh.size() <= 1: 

793 raise ValueError( 

794 f"Hybrid wrapper {wrapper_name!r} requires an active CP mesh" 

795 ) 

796 if isinstance(ulysses_degree, bool) or not isinstance(ulysses_degree, int): 

797 raise TypeError( 

798 f"Hybrid wrapper {wrapper_name!r} requires integer " 

799 f"ulysses_degree, got {type(ulysses_degree).__name__}" 

800 ) 

801 cp_size = cp_mesh.size() 

802 if not 1 < ulysses_degree < cp_size: 

803 raise ValueError( 

804 f"Hybrid wrapper {wrapper_name!r} requires 1 < ulysses_degree " 

805 f"< cp_size, got ulysses_degree={ulysses_degree}, cp_size={cp_size}" 

806 ) 

807 if cp_size % ulysses_degree: 

808 raise ValueError( 

809 f"cp_size ({cp_size}) must be divisible by ulysses_degree " 

810 f"({ulysses_degree})" 

811 ) 

812 

813 

814def _apply_qkv_hybrid_wrapper( 

815 target_module, cp_mesh, ulysses_degree, wrapper_name, *, 

816 prepare_sdpa_mask: bool): 

817 """Replace a separated QKV forward with local-tensor Hybrid CP.""" 

818 _validate_hybrid_config(cp_mesh, ulysses_degree, wrapper_name) 

819 original_forward = target_module.forward 

820 

821 @functools.wraps(original_forward) 

822 def cp_forward(q: Any, k: Any, v: Any, **kwargs: Any) -> Any: 

823 """Run the original QKV forward through local-tensor Hybrid CP.""" 

824 def attention_call( 

825 query: Any, key: Any, value: Any, 

826 call_kwargs: dict[str, Any]) -> Any: 

827 """Invoke the original QKV forward with transformed tensors.""" 

828 return original_forward(query, key, value, **call_kwargs) 

829 

830 if prepare_sdpa_mask: 

831 kwargs = _prepare_hybrid_sdpa_kwargs( 

832 kwargs, q, k, cp_mesh, ulysses_degree 

833 ) 

834 return hybrid_cp_attention( 

835 attention_call, 

836 q, 

837 k, 

838 v, 

839 kwargs, 

840 cp_mesh, 

841 ulysses_degree, 

842 ) 

843 

844 return cp_forward 

845 

846 

847@inner_wrapper 

848def sdpa_qkv_hybrid_cp_wrapper( 

849 target_module: Module, mesh: Any, tp_mesh: Any, cp_mesh: Any, 

850 ep_mesh: Any, ulysses_degree: int) -> Callable: 

851 """Apply Hybrid CP to a separated SDPA ``forward(query, key, value, ...)``.""" 

852 del mesh, tp_mesh, ep_mesh 

853 return _apply_qkv_hybrid_wrapper( 

854 target_module, 

855 cp_mesh, 

856 ulysses_degree, 

857 "sdpa_qkv_hybrid", 

858 prepare_sdpa_mask=True, 

859 ) 

860 

861 

862@inner_wrapper 

863def flex_qkv_hybrid_cp_wrapper( 

864 target_module: Module, mesh: Any, tp_mesh: Any, cp_mesh: Any, 

865 ep_mesh: Any, ulysses_degree: int) -> Callable: 

866 """Apply Hybrid CP to a separated FlexAttention QKV forward.""" 

867 del mesh, tp_mesh, ep_mesh 

868 return _apply_qkv_hybrid_wrapper( 

869 target_module, 

870 cp_mesh, 

871 ulysses_degree, 

872 "flex_qkv_hybrid", 

873 prepare_sdpa_mask=False, 

874 ) 

875 

876 

877@inner_wrapper 

878def sdpa_hf_hybrid_cp_wrapper( 

879 target_module: Module, mesh: Any, tp_mesh: Any, cp_mesh: Any, 

880 ep_mesh: Any, ulysses_degree: int) -> Callable: 

881 """Apply Hybrid CP by intercepting HF SDPA primitive calls.""" 

882 del mesh, tp_mesh, ep_mesh 

883 _validate_hybrid_config(cp_mesh, ulysses_degree, "sdpa_hf_hybrid") 

884 original_forward = target_module.forward 

885 original_sdpa = F.scaled_dot_product_attention 

886 

887 @functools.wraps(original_forward) 

888 def cp_forward(*args: Any, **kwargs: Any) -> Any: 

889 """Run the HF attention forward with temporary Hybrid SDPA interception.""" 

890 fired = {"hit": False} 

891 

892 def hybrid_sdpa( 

893 query: Any, key: Any, value: Any, 

894 **attention_kwargs: Any) -> Any: 

895 """Route one intercepted SDPA call through Hybrid CP.""" 

896 fired["hit"] = True 

897 attention_kwargs = _prepare_hybrid_sdpa_kwargs( 

898 attention_kwargs, 

899 query, 

900 key, 

901 cp_mesh, 

902 ulysses_degree, 

903 ) 

904 

905 def _call_original_sdpa(call_query, call_key, call_value, call_kwargs): 

906 return original_sdpa(