Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / context_parallel / async_context_parallel.py: 96%
273 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
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"""AsyncContextParallel: overlap projection GEMM with CP collectives.
17Supports async Pure Ulysses, async Pure Colossal AI, and Hybrid CP modes.
18Hybrid keeps the original half-async path: async Ulysses A2A + sync Colossal
19AllGather. Falls back to sync ContextParallel when q/k/v_proj are not provided.
21Forward: proj hooks launch async A2A → attn pre-hook waits Q/K/V → attn hook gathers output
22Backward: autograd backward launches async A2A → proj pre-hooks wait before GEMMs
23"""
24from functools import partial
25from typing import Optional
27from hyper_parallel.core.context_parallel.context_parallel import (
28 ContextParallel,
29 _build_hybrid_cp_mesh,
30 _compose_cp_non_cp_placements,
31 _ensure_1d,
32 _finalize_ata_output,
33 _gather_seq,
34 _gather_head_to_seq,
35 _is_cp_composed_dtensor,
36 _localize_foreign_dtensor,
37 _non_cp_dtensor_layout,
38 _pop_output_layout,
39 _to_cp_dtensor,
40)
41from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
42from hyper_parallel.core.dtensor.dtensor import DTensor
43from hyper_parallel.core.dtensor.placement_types import Shard, Replicate
44from hyper_parallel.platform import get_platform
46platform = get_platform()
47Module = platform.Module
48Tensor = platform.Tensor
51# ---------------------------------------------------------------------------
52# All-to-all helpers
53# ---------------------------------------------------------------------------
55def _detach_if_available(tensor: Tensor) -> Tensor:
56 """Detach the communication buffer when the backend tensor exposes ``detach``."""
57 detach = getattr(tensor, "detach", None)
58 return detach() if detach is not None else tensor
61def _launch_async_a2a_seq_to_head(
62 tensor: Tensor,
63 group,
64 world_size: int,
65 head_dim: int,
66) -> tuple:
67 """Launch async seq→head A2A (forward)."""
68 x = tensor.contiguous()
69 shape = list(x.shape)
70 num_heads = shape[head_dim]
71 if num_heads % world_size != 0:
72 raise ValueError(f"num_heads ({num_heads}) must be divisible by world_size ({world_size}).")
73 ndim = len(shape) + 1
74 x_perm = x.reshape(
75 shape[:head_dim] + [world_size, num_heads // world_size] + shape[head_dim + 1:]
76 ).permute(
77 [head_dim] + list(range(head_dim)) + list(range(head_dim + 1, ndim))
78 ).contiguous()
79 out_perm, work = platform.all_to_all_single(_detach_if_available(x_perm), list(x_perm.shape), group, async_op=True)
80 return work, out_perm
83def _a2a_reconstruct(out_perm: Tensor, concat_dim: int) -> Tensor:
84 """Reconstruct A2A result from raw out_perm."""
85 new_ndim = out_perm.dim()
86 chunk_in_perm = concat_dim + 1
87 recon_perm = list(range(1, chunk_in_perm)) + [0] + list(range(chunk_in_perm, new_ndim))
88 x_recon = out_perm.permute(recon_perm).contiguous()
89 shape = list(x_recon.shape)
90 merged = shape[concat_dim] * shape[concat_dim + 1]
91 return x_recon.reshape(shape[:concat_dim] + [merged] + shape[concat_dim + 2:])
94def _normalize_dim(dim: int, ndim: int) -> int:
95 """Normalize a possibly negative dimension index."""
96 return dim + ndim if dim < 0 else dim
99def _move_dim_to_front(tensor: Tensor, dim: int) -> Tensor:
100 """Move ``dim`` to the leading dimension before all-gather/reduce-scatter."""
101 dim = _normalize_dim(dim, tensor.dim())
102 if dim == 0:
103 return tensor.contiguous()
104 perm = [dim] + [i for i in range(tensor.dim()) if i != dim]
105 return tensor.permute(perm).contiguous()
108def _move_dim_from_front(tensor: Tensor, dim: int) -> Tensor:
109 """Inverse of :func:`_move_dim_to_front`."""
110 dim = _normalize_dim(dim, tensor.dim())
111 if dim == 0:
112 return tensor.contiguous()
113 perm = [dim] + [i for i in range(tensor.dim()) if i != dim]
114 inverse = [0] * len(perm)
115 for idx, value in enumerate(perm):
116 inverse[value] = idx
117 return tensor.permute(inverse).contiguous()
120def _launch_async_allgather_seq(
121 tensor: Tensor,
122 group,
123 world_size: int,
124 gather_dim: int,
125) -> tuple:
126 """Launch async all-gather along ``gather_dim``."""
127 x_perm = _move_dim_to_front(tensor.contiguous(), gather_dim)
128 output_shape = list(x_perm.shape)
129 output_shape[0] *= world_size
130 out_perm, work = platform.all_gather_single(_detach_if_available(x_perm), output_shape, group, async_op=True)
131 return work, out_perm
134def _allgather_reconstruct(out_perm: Tensor, gather_dim: int) -> Tensor:
135 """Move the leading communication buffer dimension back to ``gather_dim``.
137 The input may be a forward all-gather output buffer or a backward
138 reduce-scatter local buffer, as both are stored with ``gather_dim`` moved
139 to dimension 0 before communication.
140 """
141 return _move_dim_from_front(out_perm, gather_dim)
144# ---------------------------------------------------------------------------
145# Tensor helpers
146# ---------------------------------------------------------------------------
148def _to_local(tensor: Tensor) -> Tensor:
149 """Return the local Tensor from a DTensor, or the tensor itself."""
150 return tensor.to_local() if isinstance(tensor, DTensor) else tensor
153def _to_cp_local(tensor: Tensor, submesh: DeviceMesh, seq_dim: int) -> Tensor:
154 """Return the local tensor CP async collectives should consume."""
155 tensor = _localize_foreign_dtensor(tensor, submesh, seq_dim)
156 return tensor.to_local() if isinstance(tensor, DTensor) else tensor
159def _cp_execution_layout_from_input(tensor: Tensor, cp_mesh: DeviceMesh, cp_placements, seq_dim: int) -> tuple:
160 """Return the DTensor layout to use after a raw async CP collective."""
161 cp_placements = tuple(cp_placements)
162 layout = _non_cp_dtensor_layout(tensor, cp_mesh, seq_dim)
163 if layout is None:
164 return cp_mesh, cp_placements
165 non_cp_mesh, non_cp_placements, composed_mesh = layout
166 return composed_mesh, _compose_cp_non_cp_placements(
167 composed_mesh,
168 cp_mesh,
169 cp_placements,
170 non_cp_mesh,
171 non_cp_placements,
172 )
175def _make_async_cp_slot(work, out_perm, tensor: Tensor, cp_mesh: DeviceMesh, cp_placements, seq_dim: int) -> dict:
176 """Store an async handle together with the layout needed after wait."""
177 return {
178 "work": work,
179 "out_perm": out_perm,
180 "layout": _cp_execution_layout_from_input(tensor, cp_mesh, cp_placements, seq_dim),
181 }
184def _slot_comm(slot):
185 """Return ``(work, out_perm)`` from a new metadata slot or a legacy tuple."""
186 if isinstance(slot, dict):
187 return slot["work"], slot["out_perm"]
188 return slot
191def _slot_layout(slot, fallback_mesh: DeviceMesh, fallback_placements) -> tuple:
192 """Return the post-collective layout recorded in *slot*."""
193 if isinstance(slot, dict) and "layout" in slot:
194 return slot["layout"]
195 return fallback_mesh, tuple(fallback_placements)
198def _wrap_async_cp_result(local: Tensor, slot, fallback_mesh: DeviceMesh, fallback_placements) -> "DTensor":
199 """Wrap a waited async result with the layout captured before communication."""
200 mesh, placements = _slot_layout(slot, fallback_mesh, fallback_placements)
201 return DTensor.from_local(local, mesh, placements)
204# ---------------------------------------------------------------------------
205# AsyncContextParallel
206# ---------------------------------------------------------------------------
208class AsyncContextParallel(ContextParallel):
209 """Context Parallel with projection–collective compute overlap.
211 Requires ``q_proj``, ``k_proj``, ``v_proj`` in :meth:`apply`; otherwise
212 falls back to synchronous :class:`ContextParallel`.
214 Pure Colossal AI (``ulysses_degree=1``) overlaps K/V AllGather when the
215 backend exposes async all-gather handles. Hybrid CP overlaps Ulysses A2A and
216 keeps the Colossal K/V AllGather synchronous for backward-ordering safety.
218 Args:
219 seq_dim: Sequence dimension (1=BSHD, 2=BNSD).
220 head_dim: Head dimension (2=BSHD, 1=BNSD).
221 ulysses_degree: Ulysses sub-mesh size (see :class:`ContextParallel`).
222 qkv_indices: Positional indices of (Q, K, V) in attention forward.
223 qkv_kwarg_names: Keyword names for (Q, K, V).
224 use_local_output: Return local tensors after CP when True. When False,
225 keep CP DTensor outputs, or drop the CP axis from
226 composed CP+TP outputs and keep the non-CP layout.
227 load_balance: Load-balance flag forwarded to base class.
228 """
230 def __init__(
231 self,
232 seq_dim: int = 1,
233 head_dim: int = 2,
234 ulysses_degree: Optional[int] = None,
235 qkv_indices: tuple = (0, 1, 2),
236 qkv_kwarg_names: tuple = (),
237 use_local_output: bool = False,
238 load_balance: bool = False,
239 ):
240 super().__init__(
241 seq_dim=seq_dim,
242 head_dim=head_dim,
243 ulysses_degree=ulysses_degree,
244 qkv_indices=qkv_indices,
245 qkv_kwarg_names=qkv_kwarg_names,
246 use_local_output=use_local_output,
247 load_balance=load_balance,
248 )
250 # ------------------------------------------------------------------
251 # Public entry point
252 # ------------------------------------------------------------------
254 def apply( # pylint: disable=arguments-differ
255 self,
256 module: Module,
257 device_mesh: DeviceMesh,
258 q_proj: Optional[Module] = None,
259 k_proj: Optional[Module] = None,
260 v_proj: Optional[Module] = None,
261 ) -> Module:
262 """Register async-overlap hooks and return *module*.
264 Falls back to synchronous :class:`ContextParallel` if any of
265 ``q/k/v_proj`` is ``None``.
267 Args:
268 module: Core-attention submodule.
269 device_mesh: CP device mesh (1-D or 2-D).
270 q_proj: The last module in the Q path whose output is passed
271 directly to the attention module as Q. Its forward
272 post-hook launches the async Q all-to-all. There
273 must be **no** intermediate ops (view, transpose, …)
274 between this module and attention; such ops would be
275 bypassed by the pre-hook substitution and could cause
276 shape mismatches. For models with QK normalization
277 applied right before attention, pass ``qk_norm_q``
278 here instead of the raw projection.
279 k_proj: Same semantics as ``q_proj``, for the K path. Pass
280 ``qk_norm_k`` when the model applies QK-Norm before
281 attention.
282 v_proj: Value projection module (no norm variant needed).
283 """
284 if q_proj is None or k_proj is None or v_proj is None:
285 return super().apply(module, device_mesh)
287 cp_size = device_mesh.mesh.numel()
288 ds = self.ulysses_degree if self.ulysses_degree is not None else cp_size
289 if cp_size % ds != 0:
290 raise ValueError(
291 f"cp_size ({cp_size}) must be divisible by ulysses_degree ({ds})."
292 )
293 co = cp_size // ds
295 if ds == 1:
296 if self.load_balance:
297 return super().apply(module, device_mesh)
298 return self._apply_colossal_async(module, device_mesh, cp_size, k_proj, v_proj)
300 return self._apply_a2a_async(module, device_mesh, ds, co, q_proj, k_proj, v_proj)
302 def _apply_colossal_async(
303 self,
304 module: Module,
305 device_mesh: DeviceMesh,
306 cp_size: int,
307 k_proj: Module,
308 v_proj: Module,
309 ) -> Module:
310 """Register Pure Colossal async K/V AllGather hooks."""
311 co_submesh = _ensure_1d(device_mesh)
312 group = co_submesh.get_group()
313 fwd_ag_slots = {"k": None, "v": None}
314 bwd_ag_slots = {"k": [], "v": []}
315 self._register_ag_proj_hooks(
316 k_proj,
317 v_proj,
318 submesh=co_submesh,
319 group=group,
320 world_size=cp_size,
321 fwd_slots=fwd_ag_slots,
322 bwd_slots=bwd_ag_slots,
323 )
324 platform.register_forward_pre_hook(
325 module,
326 partial(
327 self._attn_pre_hook_colossal,
328 co_submesh=co_submesh,
329 group=group,
330 world_size=cp_size,
331 fwd_slots=fwd_ag_slots,
332 bwd_slots=bwd_ag_slots,
333 ),
334 with_kwargs=True,
335 )
336 module.register_forward_hook(
337 partial(self._post_hook_colossal, co_submesh=co_submesh)
338 )
339 return module
341 def _apply_a2a_async( # pylint: disable=too-many-arguments
342 self,
343 module: Module,
344 device_mesh: DeviceMesh,
345 ds: int,
346 co: int,
347 q_proj: Module,
348 k_proj: Module,
349 v_proj: Module,
350 ) -> Module:
351 """Register Pure Ulysses or Hybrid async A2A hooks."""
352 fwd_slots = {"q": None, "k": None, "v": None}
353 bwd_slots = {"q": [], "k": [], "v": []}
355 if co == 1:
356 ds_submesh = _ensure_1d(device_mesh)
357 group = ds_submesh.get_group()
358 a2a_layout_mesh = ds_submesh
359 a2a_layout_placements = (Shard(self.head_dim),)
360 pre_hook = partial(
361 self._attn_pre_hook_ulysses,
362 ds_submesh=ds_submesh,
363 group=group,
364 world_size=ds,
365 fwd_slots=fwd_slots,
366 bwd_slots=bwd_slots,
367 )
368 post_hook = partial(self._attn_post_hook_ata, ds_submesh=ds_submesh)
369 else:
370 hybrid_cp_mesh = _build_hybrid_cp_mesh(device_mesh, ds, co)
371 dim_names = hybrid_cp_mesh.mesh_dim_names
372 if dim_names is None:
373 raise ValueError("2-D mesh must have mesh_dim_names (guaranteed by _build_2d_mesh)")
374 ds_submesh = hybrid_cp_mesh[dim_names[1]]
375 group = ds_submesh.get_group()
376 a2a_layout_mesh = hybrid_cp_mesh
377 a2a_layout_placements = (Shard(self.seq_dim), Shard(self.head_dim))
378 pre_hook = partial(
379 self._attn_pre_hook_hybrid,
380 ds_submesh=ds_submesh,
381 group=group,
382 world_size=ds,
383 hybrid_cp_mesh=hybrid_cp_mesh,
384 fwd_slots=fwd_slots,
385 bwd_slots=bwd_slots,
386 )
387 post_hook = partial(
388 self._post_hook_hybrid,
389 hybrid_cp_mesh=hybrid_cp_mesh,
390 ds_submesh=ds_submesh,
391 )
393 self._register_proj_hooks(
394 q_proj,
395 k_proj,
396 v_proj,
397 submesh=ds_submesh,
398 group=group,
399 world_size=ds,
400 fwd_slots=fwd_slots,
401 bwd_slots=bwd_slots,
402 layout_mesh=a2a_layout_mesh,
403 layout_placements=a2a_layout_placements,
404 )
405 platform.register_forward_pre_hook(
406 module,
407 pre_hook,
408 with_kwargs=True,
409 )
410 module.register_forward_hook(post_hook)
411 return module
413 # ------------------------------------------------------------------
414 # Shared: projection hooks registration
415 # ------------------------------------------------------------------
417 def _register_proj_hooks(
418 self,
419 q_proj,
420 k_proj,
421 v_proj,
422 submesh,
423 group,
424 world_size,
425 fwd_slots,
426 bwd_slots,
427 layout_mesh,
428 layout_placements,
429 ):
430 """Register forward and backward hooks on all three projection modules."""
431 for key, proj in [("q", q_proj), ("k", k_proj), ("v", v_proj)]:
432 proj.register_forward_hook(
433 partial(self._proj_post_hook, key=key, submesh=submesh, group=group, world_size=world_size,
434 fwd_slots=fwd_slots, layout_mesh=layout_mesh, layout_placements=layout_placements)
435 )
436 platform.register_full_backward_pre_hook(
437 proj,
438 partial(self._proj_bwd_pre_hook, bwd_slot=bwd_slots[key])
439 )
441 def _proj_post_hook( # pylint: disable=unused-argument,too-many-arguments
442 self, module, inputs, output, key, submesh, group, world_size, fwd_slots, layout_mesh, layout_placements
443 ):
444 """Launch async seq→head A2A after projection; return original output unchanged."""
445 tensor = _to_cp_local(output, submesh, self.seq_dim)
446 work, out_perm = _launch_async_a2a_seq_to_head(
447 tensor, group, world_size, self.head_dim
448 )
449 fwd_slots[key] = _make_async_cp_slot(
450 work,
451 out_perm,
452 output,
453 layout_mesh,
454 layout_placements,
455 self.seq_dim,
456 )
457 return output
459 def _register_ag_proj_hooks(self, k_proj, v_proj, submesh, group, world_size, fwd_slots, bwd_slots):
460 """Register async AllGather hooks for K/V projection modules."""
461 for key, proj in [("k", k_proj), ("v", v_proj)]:
462 proj.register_forward_hook(
463 partial(self._proj_ag_post_hook, key=key, submesh=submesh, group=group, world_size=world_size,
464 fwd_slots=fwd_slots)
465 )
466 platform.register_full_backward_pre_hook(
467 proj,
468 partial(self._proj_ag_bwd_pre_hook, bwd_slot=bwd_slots[key])
469 )
471 def _proj_ag_post_hook( # pylint: disable=unused-argument,too-many-arguments
472 self, module, inputs, output, key, submesh, group, world_size, fwd_slots
473 ):
474 """Launch async K/V AllGather after projection; return original output."""
475 tensor = _to_cp_local(output, submesh, self.seq_dim)
476 work, out_perm = _launch_async_allgather_seq(
477 tensor, group, world_size, self.seq_dim
478 )
479 fwd_slots[key] = _make_async_cp_slot(
480 work,
481 out_perm,
482 output,
483 submesh,
484 (Replicate(),),
485 self.seq_dim,
486 )
487 return output
489 def _get_qkv_value(self, args, kwargs, qkv_pos: int):
490 """Return Q/K/V value from positional args or configured kwargs."""
491 idx = self.qkv_indices[qkv_pos]
492 if idx < len(args):
493 return args[idx]
494 if qkv_pos < len(self.qkv_kwarg_names):
495 name = self.qkv_kwarg_names[qkv_pos]
496 if name in kwargs:
497 return kwargs[name]
498 return None
500 def _set_qkv_value(self, args, kwargs, qkv_pos: int, value):
501 """Set Q/K/V value in positional args or configured kwargs."""
502 idx = self.qkv_indices[qkv_pos]
503 if idx < len(args):
504 args[idx] = value
505 return
506 if qkv_pos < len(self.qkv_kwarg_names):
507 name = self.qkv_kwarg_names[qkv_pos]
508 if name in kwargs:
509 kwargs[name] = value
511 # ------------------------------------------------------------------
512 # Internal: wait for pre-launched communication handles
513 # ------------------------------------------------------------------
515 def _wait_a2a(self, tensor, group, world_size, fwd_slots, key, bwd_slot):
516 """Wait for pre-launched A2A; returns head-scattered tensor (differentiable)."""
517 slot = fwd_slots[key]
518 work, out_perm = _slot_comm(slot)
519 fwd_slots[key] = None
520 return platform.differentiable_async_a2a_wait(
521 tensor, work, out_perm, group, world_size,
522 self.seq_dim, self.head_dim, # concat_dim=seq_dim, split_dim=head_dim
523 bwd_slot,
524 )
526 def _wait_a2a_dtensor(
527 self,
528 tensor,
529 group,
530 world_size,
531 fwd_slots,
532 key,
533 bwd_slot,
534 fallback_mesh,
535 fallback_placements,
536 ):
537 """Wait for pre-launched A2A and wrap the result with captured CP+non-CP layout."""
538 slot = fwd_slots[key]
539 local = self._wait_a2a(
540 tensor,
541 group,
542 world_size,
543 fwd_slots,
544 key,
545 bwd_slot,
546 )
547 return _wrap_async_cp_result(local, slot, fallback_mesh, fallback_placements)
549 def _wait_allgather(self, tensor, group, world_size, work, out_perm, bwd_slot=None):
550 """Wait for pre-launched AllGather and return gathered tensor."""
551 return platform.differentiable_async_allgather_wait(
552 tensor,
553 work,
554 out_perm,
555 group,
556 world_size,
557 self.seq_dim,
558 bwd_slot,
559 )
561 # ------------------------------------------------------------------
562 # QKV transform helpers
563 # ------------------------------------------------------------------
565 def _apply_qkv_transforms(
566 self,
567 args,
568 kwargs,
569 submesh,
570 transform_q,
571 transform_k,
572 transform_v,
573 ):
574 """Apply transforms to Q/K/V and return modified ``(args, kwargs)``.
576 Each transform receives ``(original, local)``. The local tensor is the
577 raw communication buffer, while the original value carries any non-CP
578 DTensor layout that must be restored after the async wait.
579 """
580 new_args = list(args)
581 new_kwargs = dict(kwargs)
583 transforms = (transform_q, transform_k, transform_v)
584 for pos, transform in enumerate(transforms):
585 value = self._get_qkv_value(new_args, new_kwargs, pos)
586 if value is None:
587 continue
588 self._set_qkv_value(
589 new_args,
590 new_kwargs,
591 pos,
592 transform(value, _to_cp_local(value, submesh, self.seq_dim)),
593 )
594 return tuple(new_args), new_kwargs
596 # ------------------------------------------------------------------
597 # Attention pre-hooks
598 # ------------------------------------------------------------------
600 def _attn_pre_hook_ulysses( # pylint: disable=unused-argument,too-many-arguments
601 self, module, args, kwargs, group, world_size, fwd_slots, bwd_slots, ds_submesh=None
602 ):
603 """Wait Q/K/V A2A and return head-scattered args/kwargs."""
604 self._record_q_output_tp_layout(module, args, kwargs, ds_submesh)
606 def make_a2a(key):
607 def transform(value, local_value): # pylint: disable=unused-argument
608 if ds_submesh is None:
609 return self._wait_a2a(
610 local_value,
611 group,
612 world_size,
613 fwd_slots,
614 key,
615 bwd_slots[key],
616 )
617 return self._wait_a2a_dtensor(
618 local_value,
619 group,
620 world_size,
621 fwd_slots,
622 key,
623 bwd_slots[key],
624 ds_submesh,
625 (Shard(self.head_dim),),
626 )
627 return transform
629 return self._apply_qkv_transforms(
630 args,
631 kwargs,
632 ds_submesh,
633 transform_q=make_a2a("q"),
634 transform_k=make_a2a("k"),
635 transform_v=make_a2a("v"),
636 )
638 def _ulysses_cp_layout_from_input(self, t, ds_submesh):
639 """Return the pure-Ulysses execution mesh and placements for one input tensor."""
640 cp_placements = (Shard(self.head_dim),)
641 layout = _non_cp_dtensor_layout(t, ds_submesh, self.seq_dim)
642 if layout is None:
643 return ds_submesh, cp_placements
644 non_cp_mesh, non_cp_placements, composed_mesh = layout
645 return composed_mesh, _compose_cp_non_cp_placements(
646 composed_mesh,
647 ds_submesh,
648 cp_placements,
649 non_cp_mesh,
650 non_cp_placements,
651 )
653 def _attn_pre_hook_colossal( # pylint: disable=unused-argument,too-many-arguments
654 self, module, args, kwargs, co_submesh, group, world_size, fwd_slots, bwd_slots
655 ):
656 """Wait K/V AllGather, wrap Q as local shard and return CP-ready args/kwargs."""
657 new_args = list(args)
658 new_kwargs = dict(kwargs)
660 q_value = self._get_qkv_value(new_args, new_kwargs, 0)
661 self._record_q_output_tp_layout(module, new_args, new_kwargs, co_submesh)
662 if q_value is not None:
663 self._set_qkv_value(
664 new_args,
665 new_kwargs,
666 0,
667 _to_cp_dtensor(
668 q_value,
669 co_submesh,
670 (Shard(self.seq_dim),),
671 (Shard(self.seq_dim),),
672 self.seq_dim,
673 ),
674 )
676 for pos, key in enumerate(("k", "v"), start=1):
677 value = self._get_qkv_value(new_args, new_kwargs, pos)
678 if value is None:
679 continue
680 local = _to_cp_local(value, co_submesh, self.seq_dim)
681 slot = fwd_slots[key]
682 work, out_perm = _slot_comm(slot)
683 fwd_slots[key] = None
684 gathered = self._wait_allgather(
685 local,
686 group,
687 world_size,
688 work,
689 out_perm,
690 bwd_slots[key],
691 )
692 if isinstance(value, DTensor):
693 layout = _non_cp_dtensor_layout(value, co_submesh, self.seq_dim)
694 if layout is not None:
695 non_cp_mesh, non_cp_placements, composed_mesh = layout
696 composed_placements = _compose_cp_non_cp_placements(
697 composed_mesh,
698 co_submesh,
699 (Replicate(),),
700 non_cp_mesh,
701 non_cp_placements,
702 )
703 self._set_qkv_value(
704 new_args,
705 new_kwargs,
706 pos,
707 DTensor.from_local(
708 gathered.to_local() if isinstance(gathered, DTensor) else gathered,
709 composed_mesh,
710 composed_placements,
711 ),
712 )
713 continue
714 self._set_qkv_value(
715 new_args,
716 new_kwargs,
717 pos,
718 DTensor.from_local(
719 _to_local(gathered),
720 co_submesh,
721 (Replicate(),),
722 ),
723 )
725 return tuple(new_args), new_kwargs
727 def _attn_pre_hook_hybrid( # pylint: disable=too-many-locals,too-many-arguments
728 self, module, args, kwargs, group, world_size, hybrid_cp_mesh, # pylint: disable=unused-argument
729 fwd_slots, bwd_slots, ds_submesh=None
730 ):
731 """Wait Ulysses A2A for Q/K/V, then synchronously gather K/V on the co-submesh."""
732 new_args = list(args)
733 new_kwargs = dict(kwargs)
735 co_submesh = hybrid_cp_mesh[hybrid_cp_mesh.mesh_dim_names[0]]
736 if ds_submesh is None:
737 ds_submesh = hybrid_cp_mesh[hybrid_cp_mesh.mesh_dim_names[1]]
738 self._record_q_output_tp_layout(module, new_args, new_kwargs, hybrid_cp_mesh)
740 q_value = self._get_qkv_value(new_args, new_kwargs, 0)
741 if q_value is not None:
742 self._set_qkv_value(
743 new_args,
744 new_kwargs,
745 0,
746 self._wait_a2a_dtensor(
747 _to_cp_local(q_value, ds_submesh, self.seq_dim),
748 group,
749 world_size,
750 fwd_slots,
751 "q",
752 bwd_slots["q"],
753 hybrid_cp_mesh,
754 (Shard(self.seq_dim), Shard(self.head_dim)),
755 ),
756 )
758 k_value = self._get_qkv_value(new_args, new_kwargs, 1)
759 if k_value is not None:
760 k_ul = self._wait_a2a_dtensor(
761 _to_cp_local(k_value, ds_submesh, self.seq_dim),
762 group,
763 world_size,
764 fwd_slots,
765 "k",
766 bwd_slots["k"],
767 hybrid_cp_mesh,
768 (Shard(self.seq_dim), Shard(self.head_dim)),
769 )
770 k_full = _gather_seq(k_ul, co_submesh, self.seq_dim)
771 k_mesh, k_placements = self._hybrid_kv_gather_placements(
772 k_value,
773 hybrid_cp_mesh,
774 )
775 self._set_qkv_value(
776 new_args,
777 new_kwargs,
778 1,
779 DTensor.from_local(
780 _to_local(k_full),
781 k_mesh,
782 k_placements,
783 ),
784 )
786 v_value = self._get_qkv_value(new_args, new_kwargs, 2)
787 if v_value is not None:
788 v_ul = self._wait_a2a_dtensor(
789 _to_cp_local(v_value, ds_submesh, self.seq_dim),
790 group,
791 world_size,
792 fwd_slots,
793 "v",
794 bwd_slots["v"],
795 hybrid_cp_mesh,
796 (Shard(self.seq_dim), Shard(self.head_dim)),
797 )
798 v_full = _gather_seq(v_ul, co_submesh, self.seq_dim)
799 v_mesh, v_placements = self._hybrid_kv_gather_placements(
800 v_value,
801 hybrid_cp_mesh,
802 )
803 self._set_qkv_value(
804 new_args,
805 new_kwargs,
806 2,
807 DTensor.from_local(
808 _to_local(v_full),
809 v_mesh,
810 v_placements,
811 ),
812 )
814 return tuple(new_args), new_kwargs
816 # ------------------------------------------------------------------
817 # Attention post-hook (Ulysses and Hybrid share the same reverse ATA)
818 # ------------------------------------------------------------------
820 def _attn_post_hook_ata(self, module, args, output, ds_submesh): # pylint: disable=unused-argument
821 """Reverse head→seq gather on ds_submesh; returns local tensor."""
822 output_layout = _pop_output_layout(module)
824 def _process(output_item):
825 if isinstance(output_item, (Tensor, DTensor)):
826 if isinstance(output_item, DTensor) and not _is_cp_composed_dtensor(output_item, ds_submesh):
827 output_item = output_item.to_local()
828 seq_dtensor = _gather_head_to_seq(
829 output_item, ds_submesh, self.seq_dim, self.head_dim
830 )
831 return _finalize_ata_output(
832 seq_dtensor, output_layout, ds_submesh, self.seq_dim, self.use_local_output
833 )
834 return output_item
836 if isinstance(output, (tuple, list)):
837 return type(output)(_process(item) for item in output)
838 return _process(output)
840 # ------------------------------------------------------------------
841 # Backward: wait A2A handle (launched by autograd) before proj GEMM
842 # ------------------------------------------------------------------
844 def _proj_bwd_pre_hook(self, module, grad_output, bwd_slot): # pylint: disable=unused-argument
845 """Wait backward A2A just before proj GEMM; replace grad with seq-form.
847 The async head→seq A2A is launched inside _TorchAsyncA2AFunction.backward
848 and appended to ``bwd_slot``. Waiting here lets the A2A overlap with the
849 preceding proj GEMM.
850 """
851 work, out_perm = bwd_slot.pop()
852 work.wait()
853 d_seq = _a2a_reconstruct(out_perm, self.head_dim)
854 return (d_seq,) + grad_output[1:] if isinstance(grad_output, tuple) else (d_seq,)
856 def _proj_ag_bwd_pre_hook(self, module, grad_output, bwd_slot): # pylint: disable=unused-argument
857 """Wait backward reduce-scatter just before K/V projection GEMM."""
858 work, out_perm, gather_dim = bwd_slot.pop()
859 work.wait()
860 d_local = _allgather_reconstruct(out_perm, gather_dim)
861 return (d_local,) + grad_output[1:] if isinstance(grad_output, tuple) else (d_local,)