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