Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / wrappers.py: 43%
438 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"""context_parallel.wrappers: built-in CP-aware inner attention wrappers (public, explicit injection).
17Since the explicit-injection rework the framework NEVER picks a CP wrapper
18automatically: the built-ins below are referenced explicitly — by
19registry name (``spec.inner_wrapper="sdpa_hf"``), by callable, or by a YAML
20Target pointing at one of these functions:
22.. code-block:: yaml
24 plan_overrides:
25 - match: "*.self_attn"
26 when: cp
27 inner_wrapper:
28 _target_: hyper_parallel.distributed.context_parallel.wrappers.sdpa_hf_cp_wrapper
30Wrapper contract (also the Target factory contract — at apply time the
31Target is built with its declared context, filled by name)::
33 @inner_wrapper
34 fn(target_module, mesh, tp_mesh, cp_mesh, ep_mesh, <configured params...>)
36The ``@inner_wrapper`` decorator is MANDATORY (injection discipline,
37see recipe_spec.py): the mesh family (``mesh``/``tp_mesh``/``cp_mesh``/
38``ep_mesh``) plus the anchor ``target_module`` are REQUIRED context
39params, ALL filled by the framework at apply time (None for inactive
40axes) — the user just uses them.
41Undecorated wrappers fail fast in the resolution chain; *args/**kwargs
42are forbidden. In-repo wrappers RETURN the replacement — a forward
43callable or _ForwardRewriteRequest(s) — and the forward rewriter
44validates, installs and rolls back (05 §15.2.3: module.forward is
45assigned only by distributed/_builder/forward_rewriter.py); external
46wrappers may still replace ``target_module.forward`` in place and return
47None. The replaced forward must accept the original forward's params
48(validated at apply time).
50- ``sdpa_qkv_cp_wrapper``: NeMo convention ``forward(q, k, v, ...)`` — explicit
51 K/V all-gather + D-04 offset-aware causal mask (fixes the is_causal
52 alignment error under CP);
53- ``sdpa_hf_cp_wrapper``: HF convention ``forward(hidden_states, ...)`` —
54 primitive interception of ``F.scaled_dot_product_attention`` (reuses HF
55 projections/RoPE), with misfire detection (raises if no call is
56 intercepted);
57- ``flex_qkv_cp_wrapper`` / ``flex_hf_cp_wrapper``: the two isomorphic
58 FlexAttention entries (block_mask must be built for the global kv length).
59- ``*_ulysses_cp_wrapper``: Pure Ulysses variants that exchange Q/K/V from
60 sequence shards to head shards before attention and restore the local
61 sequence shard afterward.
62- ``*_hybrid_cp_wrapper``: local-tensor Hybrid variants that run Ulysses
63 all-to-all inside each subgroup and K/V all-gather across the complementary
64 Colossal subgroup.
65- ``sdpa_*_load_balance_cp_wrapper``: local-tensor Colossal Head-Tail variants
66 for Q exchange, K/V all-gather, dual SDPA execution, output restoration,
67 and backward communication.
69Users may register their own named schemes::
71 INNER_WRAPPER_REGISTRY["my_flash"] = my_fn
73after which ``spec.inner_wrapper="my_flash"`` references it by name.
75The ``is_*_attention`` helpers are public utilities for authors of custom
76wrappers (choosing a scheme programmatically); the framework itself no
77longer dispatches heuristically on them.
78"""
80import functools
81import importlib
82import inspect
83import logging
84from typing import Any, Callable
86import torch # pylint: disable=forbidden-backend-import
87import torch.nn.functional as F
88from torch.nn import Module
90from hyper_parallel.distributed._builder.forward_rewriter import (
91 _ForwardRewriteRequest,
92)
93from hyper_parallel.distributed.context_parallel.collectives import (
94 _build_hybrid_cp_submeshes,
95 _ULYSSES_WRAPPED_FLAG,
96 _UlyssesContext,
97 _slice_sequence,
98 _validate_ulysses_head_counts,
99 flex_cp_allgather,
100 hybrid_cp_attention,
101 ulysses_head_to_seq,
102 ulysses_seq_to_head,
103)
104from hyper_parallel.distributed.context_parallel.attention import (
105 _cp_offset_causal_mask,
106 _dsa_cp_alltoall,
107 _mla_cp_alltoall,
108 _mome_cp_halo_exchange,
109 head_tail_load_balance_attention,
110)
111from hyper_parallel.distributed.recipe_spec import (
112 inner_wrapper,
113)
115logger = logging.getLogger(__name__)
118# ────────────────────────────────────────────────────────────────────────────
119# Style-detection helpers (public utilities for custom wrapper authors)
120# ────────────────────────────────────────────────────────────────────────────
122def _attn_implementation(module):
123 cfg = getattr(module, "config", None)
124 impl = getattr(cfg, "_attn_implementation", None)
125 if impl is None and isinstance(cfg, dict):
126 impl = cfg.get("attn_implementation")
127 return impl
130def is_sdpa_attention(module: Any) -> bool:
131 """True when the module is configured for or named as SDPA attention."""
132 impl = _attn_implementation(module)
133 return (impl == "sdpa") or ("SdpaAttention" in type(module).__name__)
136def is_flex_attention(module: Any) -> bool:
137 """True when the module is configured for or named as FlexAttention."""
138 impl = _attn_implementation(module)
139 return (impl == "flex_attention") or ("FlexAttention" in type(module).__name__)
142def is_hf_style_attention(module: Any) -> bool:
143 """HF style (forward(hidden_states,...), projections inside forward) -> the primitive-interception path."""
144 has_proj = (hasattr(module, "q_proj") and hasattr(module, "k_proj")
145 and hasattr(module, "v_proj"))
146 if not has_proj:
147 return False
148 try:
149 sig = inspect.signature(module.forward)
150 first_param = next(iter(sig.parameters.values()), None)
151 return first_param is not None and first_param.name == "hidden_states"
152 except (ValueError, TypeError):
153 return type(module).__name__.endswith("Attention")
156# ────────────────────────────────────────────────────────────────────────────
157# CP-aware SDPA core (K/V all-gather + D-04 offset causal mask)
158# ────────────────────────────────────────────────────────────────────────────
160def _prepare_cp_sdpa_kwargs(
161 kwargs: dict[str, Any], *, q_len: int, kv_len: int,
162 query_offset: int, device: Any) -> dict[str, Any]:
163 """Adapt a complete SDPA mask to one CP-local query interval."""
164 call_kwargs = dict(kwargs)
165 attention_mask = call_kwargs.get("attn_mask")
166 if attention_mask is not None:
167 if not isinstance(attention_mask, torch.Tensor):
168 raise TypeError(
169 "CP SDPA attn_mask must be a Tensor or None, got "
170 f"{type(attention_mask).__name__}"
171 )
172 if attention_mask.ndim < 2:
173 raise ValueError(
174 "CP SDPA attn_mask must expose query and KV dimensions, got "
175 f"shape {tuple(attention_mask.shape)}"
176 )
177 if attention_mask.shape[-1] != kv_len:
178 raise ValueError(
179 "CP SDPA attn_mask must cover the complete KV sequence: "
180 f"mask KV length={attention_mask.shape[-1]}, expected {kv_len}"
181 )
182 mask_q_len = attention_mask.shape[-2]
183 if mask_q_len != q_len:
184 if mask_q_len < query_offset + q_len:
185 raise ValueError(
186 "CP SDPA attn_mask does not cover the required query range "
187 f"[{query_offset}, {query_offset + q_len}); mask query "
188 f"length={mask_q_len}"
189 )
190 attention_mask = attention_mask.narrow(
191 -2, query_offset, q_len
192 )
193 call_kwargs["attn_mask"] = attention_mask
194 if call_kwargs.get("is_causal"):
195 call_kwargs["is_causal"] = False
196 return call_kwargs
198 if call_kwargs.get("is_causal"):
199 call_kwargs["is_causal"] = False
200 call_kwargs["attn_mask"] = _cp_offset_causal_mask(
201 q_len, kv_len, query_offset, device
202 )
203 return call_kwargs
206def _prepare_head_tail_sdpa_kwargs(
207 kwargs: dict[str, Any], query: Any, key: Any,
208 cp_mesh: Any) -> tuple[dict[str, Any], dict[str, Any]]:
209 """Prepare complete-mask query rows for both Head-Tail attention calls."""
210 local_q_len = query.shape[2]
211 if local_q_len % 2:
212 raise ValueError(
213 "Head-Tail load balance requires an even local Q sequence "
214 f"length, got {local_q_len}"
215 )
216 half_q_len = local_q_len // 2
217 local_rank = cp_mesh.get_local_rank()
218 peer_rank = cp_mesh.size() - 1 - local_rank
219 global_kv_len = key.shape[2] * cp_mesh.size()
220 keep_kwargs = _prepare_cp_sdpa_kwargs(
221 kwargs,
222 q_len=half_q_len,
223 kv_len=global_kv_len,
224 query_offset=local_rank * local_q_len,
225 device=query.device,
226 )
227 peer_kwargs = _prepare_cp_sdpa_kwargs(
228 kwargs,
229 q_len=half_q_len,
230 kv_len=global_kv_len,
231 query_offset=peer_rank * local_q_len + half_q_len,
232 device=query.device,
233 )
234 return keep_kwargs, peer_kwargs
237def _prepare_hybrid_sdpa_kwargs(
238 kwargs: dict[str, Any], query: Any, key: Any,
239 cp_mesh: Any, ulysses_degree: int) -> dict[str, Any]:
240 """Prepare complete-mask query rows for Hybrid attention."""
241 _, colossal_mesh = _build_hybrid_cp_submeshes(cp_mesh, ulysses_degree)
242 hybrid_q_len = query.shape[2] * ulysses_degree
243 return _prepare_cp_sdpa_kwargs(
244 kwargs,
245 q_len=hybrid_q_len,
246 kv_len=key.shape[2] * cp_mesh.size(),
247 query_offset=colossal_mesh.get_local_rank() * hybrid_q_len,
248 device=query.device,
249 )
252def _cp_sdpa_call(orig_sdpa, cp_mesh, q, k, v, kwargs):
253 """CP-aware SDPA: K/V all-gather + D-04 offset-aware causal mask."""
254 cp_dim = 2 # sequence dim of the [B, N, S, H] layout
255 global_k, global_v = flex_cp_allgather(
256 k.contiguous(), v.contiguous(), cp_dim, cp_mesh)
257 if cp_mesh.size() > 1:
258 kwargs = _prepare_cp_sdpa_kwargs(
259 kwargs,
260 q_len=q.shape[cp_dim],
261 kv_len=global_k.shape[cp_dim],
262 query_offset=cp_mesh.get_local_rank() * q.shape[cp_dim],
263 device=q.device,
264 )
265 return orig_sdpa(q, global_k, global_v, **kwargs)
268def _ulysses_attention_call(
269 attention_fn, cp_mesh, query, key, value, kwargs,
270 *, seq_dim=2, head_dim=1):
271 """Run attention in the Pure Ulysses full-sequence/head-sharded layout."""
272 _validate_ulysses_head_counts(
273 query, key, value, cp_mesh.size(), "Pure Ulysses CP"
274 )
275 query, key, value = (
276 ulysses_seq_to_head(tensor, seq_dim, head_dim, cp_mesh)
277 for tensor in (query, key, value)
278 )
279 output = attention_fn(query, key, value, **kwargs)
280 return ulysses_head_to_seq(output, seq_dim, head_dim, cp_mesh)
283def _ulysses_qkv_forward(original_forward, cp_mesh, args, kwargs):
284 """Preserve a QKV forward signature while applying Pure Ulysses."""
285 signature = inspect.signature(original_forward)
286 bound = signature.bind(*args, **kwargs)
287 names = list(signature.parameters)
288 if len(names) < 3:
289 raise TypeError("Ulysses QKV wrapper requires at least three forward inputs")
290 try:
291 query, key, value = (
292 bound.arguments[names[index]] for index in range(3)
293 )
294 except KeyError as exc:
295 raise TypeError(
296 "Ulysses QKV wrapper requires query, key, and value inputs"
297 ) from exc
298 _validate_ulysses_head_counts(
299 query, key, value, cp_mesh.size(), "Pure Ulysses CP"
300 )
301 query, key, value = (
302 ulysses_seq_to_head(tensor, 2, 1, cp_mesh)
303 for tensor in (query, key, value)
304 )
305 bound.arguments[names[0]] = query
306 bound.arguments[names[1]] = key
307 bound.arguments[names[2]] = value
308 output = original_forward(*bound.args, **bound.kwargs)
309 return ulysses_head_to_seq(output, 2, 1, cp_mesh)
312def _require_ulysses_cp_mesh(cp_mesh, wrapper_name):
313 """Validate the communication context required by Pure Ulysses."""
314 if cp_mesh is None or cp_mesh.size() <= 1:
315 raise ValueError(
316 f"Ulysses wrapper {wrapper_name!r} requires an active CP mesh"
317 )
320# ────────────────────────────────────────────────────────────────────────────
321# The four built-in wrappers
322# ────────────────────────────────────────────────────────────────────────────
324@inner_wrapper
325def sdpa_qkv_cp_wrapper(
326 target_module: Module, mesh: Any, tp_mesh: Any,
327 cp_mesh: Any, ep_mesh: Any) -> Callable:
328 """NeMo/Megatron SDPA path (registry "sdpa_qkv"): explicit all-gather K/V.
330 Assumes the inner_attention.forward(q,k,v,...) signature convention.
331 **Local-only** (injection discipline): the dual-mode adapter converts
332 DTensor inputs to local and re-wraps the output — this wrapper never
333 touches DTensor. Since it wraps an inner submodule, the plan MUST
334 declare ``inner_out_src: "first_input"`` (its output layout == q's
335 layout). cp_mesh is framework-filled context; fail fast when the plan
336 has no cp axis.
337 """
339 del mesh, tp_mesh, ep_mesh
340 if cp_mesh is None:
341 raise ValueError(
342 "The in-repo CP wrapper reference implementation 'sdpa_qkv' requires an "
343 "active cp axis (the communication domain for K/V all-gather), but the "
344 "framework-filled cp_mesh is None (the current plan has no cp axis) — "
345 "without CP, use a custom @inner_wrapper wrapper instead (the "
346 "cp_mesh=None semantics are then your own responsibility), or use the "
347 "local_compute_fn channel (see "
348 "examples/distributed/perf_replacement.py)")
350 original_forward = target_module.forward
352 @functools.wraps(original_forward)
353 def cp_forward(q: Any, k: Any, v: Any, **kwargs: Any) -> Any:
354 """All-gather K/V along CP, then run the original QKV SDPA forward."""
355 return _cp_sdpa_call(original_forward, cp_mesh, q, k, v, kwargs)
357 return cp_forward
360@inner_wrapper
361def flex_qkv_cp_wrapper(
362 target_module: Module, mesh: Any, tp_mesh: Any,
363 cp_mesh: Any, ep_mesh: Any) -> Callable:
364 """NeMo/Megatron FlexAttention path (registry "flex_qkv"): explicit all-gather K/V.
366 **Local-only** (see sdpa_qkv_cp_wrapper); requires
367 ``inner_out_src: "first_input"`` in the plan. Constraint: the
368 block_mask must be built for the global kv length. cp_mesh is
369 framework-filled context; fail fast when the plan has no cp axis.
370 """
372 del mesh, tp_mesh, ep_mesh
373 if cp_mesh is None:
374 raise ValueError(
375 "The in-repo CP wrapper reference implementation 'flex_qkv' requires an "
376 "active cp axis (the communication domain for K/V all-gather), but the "
377 "framework-filled cp_mesh is None (the current plan has no cp axis) — "
378 "without CP, use a custom @inner_wrapper wrapper instead (the "
379 "cp_mesh=None semantics are then your own responsibility), or use the "
380 "local_compute_fn channel (see "
381 "examples/distributed/perf_replacement.py)")
383 original_forward = target_module.forward
385 @functools.wraps(original_forward)
386 def cp_forward(q: Any, k: Any, v: Any, **kwargs: Any) -> Any:
387 """All-gather K/V along CP, then run the original QKV Flex forward."""
388 global_k, global_v = flex_cp_allgather(
389 k.contiguous(), v.contiguous(), 2, cp_mesh)
390 return original_forward(q, global_k, global_v, **kwargs)
392 return cp_forward
395@inner_wrapper
396def sdpa_hf_cp_wrapper(
397 target_module: Module, mesh: Any, tp_mesh: Any,
398 cp_mesh: Any, ep_mesh: Any) -> Callable:
399 """HF standard SDPA path: forward(hidden_states,...) -> primitive interception (05 §4.4.2).
401 **Local-only** (see sdpa_qkv_cp_wrapper): the adapter handles DTensor
402 unwrap/param-unwrap/rewrap; the rewrap uses the boundary's declared
403 out_src (this wrapper targets the boundary module itself). The
404 primitive interception is a temporary global function replacement
405 (restored via try/finally) and is not thread-safe; it is safe under
406 single-process SPMD training (consistent with the TorchTitan CP
407 implementation).
408 Misfire detection: the wrapper may be applied to a module that does not
409 call F.sdpa -- if the primitive did not intercept a single call, the K/V
410 were not gathered (a silent numerical error), so raise a RuntimeError
411 immediately.
412 """
414 del mesh, tp_mesh, ep_mesh
415 if cp_mesh is None:
416 raise ValueError(
417 "The in-repo CP wrapper reference implementation 'sdpa_hf' requires an "
418 "active cp axis (the communication domain for K/V all-gather), but the "
419 "framework-filled cp_mesh is None (the current plan has no cp axis) — "
420 "without CP, use a custom @inner_wrapper wrapper instead (the "
421 "cp_mesh=None semantics are then your own responsibility), or use the "
422 "local_compute_fn channel (see "
423 "examples/distributed/perf_replacement.py)")
425 original_forward = target_module.forward
426 orig_sdpa = F.scaled_dot_product_attention
428 @functools.wraps(original_forward)
429 def cp_forward(hidden_states: Any, *args: Any, **kwargs: Any) -> Any:
430 """Run the HF forward with F.sdpa temporarily replaced by the CP-aware one."""
431 fired = {"hit": False}
433 def cp_aware_sdpa(q: Any, k: Any, v: Any, **kw: Any) -> Any:
434 """CP-aware SDPA replacement: all-gather K/V plus the D-04 mask."""
435 fired["hit"] = True
436 return _cp_sdpa_call(orig_sdpa, cp_mesh, q, k, v, kw)
438 F.scaled_dot_product_attention = cp_aware_sdpa
439 try:
440 out = original_forward(hidden_states, *args, **kwargs)
441 finally:
442 F.scaled_dot_product_attention = orig_sdpa
443 if not fired["hit"]:
444 raise RuntimeError(
445 f"CP wrapper 'sdpa_hf' did not intercept any "
446 f"F.scaled_dot_product_attention call on "
447 f"{type(target_module).__name__} -- the wrapper does "
448 f"not match the module implementation (K/V were not "
449 f"all-gathered; continuing would produce silent numerical "
450 f"errors). Please explicitly set inner_wrapper='sdpa_qkv' "
451 f"(the (q,k,v) convention), or provide a custom inner_wrapper "
452 f"callable")
453 return out
455 return cp_forward
458@inner_wrapper
459def flex_hf_cp_wrapper(
460 target_module: Module, mesh: Any, tp_mesh: Any,
461 cp_mesh: Any, ep_mesh: Any) -> Callable:
462 """HF standard FlexAttention path: intercept flex_attention (same structure as the SDPA path).
464 **Local-only** (see sdpa_qkv_cp_wrapper). Constraint:
465 score_mod/block_mask pass through verbatim via kwargs -- under CP,
466 kv_len changes from S/cp to S, so the block_mask must be built for the
467 **global kv length** (constructed on the full sequence in the data
468 pipeline / model side), otherwise shapes and semantics are misaligned.
469 The wrapper does not validate this.
470 Misfire detection is the same as 'sdpa_hf': if no flex_attention call is
471 intercepted, raise a RuntimeError.
472 """
474 del mesh, tp_mesh, ep_mesh
475 if cp_mesh is None:
476 raise ValueError(
477 "The in-repo CP wrapper reference implementation 'flex_hf' requires an "
478 "active cp axis (the communication domain for K/V all-gather), but the "
479 "framework-filled cp_mesh is None (the current plan has no cp axis) — "
480 "without CP, use a custom @inner_wrapper wrapper instead (the "
481 "cp_mesh=None semantics are then your own responsibility), or use the "
482 "local_compute_fn channel (see "
483 "examples/distributed/perf_replacement.py)")
485 original_forward = target_module.forward
486 # FlexAttention is optional on older supported PyTorch versions.
487 flex_attention_module = importlib.import_module(
488 "torch.nn.attention.flex_attention")
489 original_flex_attention = flex_attention_module.flex_attention
491 @functools.wraps(original_forward)
492 def cp_forward(hidden_states: Any, *args: Any, **kwargs: Any) -> Any:
493 """Run the HF forward with flex_attention temporarily replaced."""
494 fired = {"hit": False}
496 def cp_aware_flex(q: Any, k: Any, v: Any, **kw: Any) -> Any:
497 """CP-aware FlexAttention replacement: all-gather K/V first."""
498 fired["hit"] = True
499 global_k, global_v = flex_cp_allgather(
500 k.contiguous(), v.contiguous(), 2, cp_mesh)
501 return original_flex_attention(q, global_k, global_v, **kw)
503 flex_attention_module.flex_attention = cp_aware_flex
504 try:
505 out = original_forward(hidden_states, *args, **kwargs)
506 finally:
507 flex_attention_module.flex_attention = original_flex_attention
508 if not fired["hit"]:
509 raise RuntimeError(
510 f"CP wrapper 'flex_hf' did not intercept any flex_attention "
511 f"call on {type(target_module).__name__} -- the wrapper "
512 f"does not match the module implementation (K/V were "
513 f"not all-gathered; continuing would produce silent numerical "
514 f"errors). Please explicitly set inner_wrapper='flex_qkv' "
515 f"(the (q,k,v) convention), or provide a custom inner_wrapper "
516 f"callable")
517 return out
519 return cp_forward
522# ────────────────────────────────────────────────────────────────────────────
523# Colossal Head-Tail load-balance wrappers (explicit plan_overrides methods)
524# ────────────────────────────────────────────────────────────────────────────
526@inner_wrapper
527def sdpa_qkv_load_balance_cp_wrapper(
528 target_module: Module, mesh: Any, tp_mesh: Any,
529 cp_mesh: Any, ep_mesh: Any) -> Callable:
530 """Apply Colossal Head-Tail load balancing to a QKV-style SDPA module.
532 Args:
533 target_module: Attention module whose first three inputs are Q/K/V.
534 mesh: Framework-owned root mesh; unused by this wrapper.
535 tp_mesh: Framework-owned TP mesh; unused by this wrapper.
536 cp_mesh: Framework-owned CP communication mesh.
537 ep_mesh: Framework-owned EP mesh; unused by this wrapper.
538 """
539 del mesh, tp_mesh, ep_mesh
540 if cp_mesh is None or cp_mesh.size() <= 1:
541 raise ValueError(
542 "sdpa_qkv_load_balance_cp_wrapper requires an active CP mesh"
543 )
544 original_forward = target_module.forward
546 @functools.wraps(original_forward)
547 def cp_forward(q: Any, k: Any, v: Any, **kwargs: Any) -> Any:
548 """Route local Q/K/V tensors through Head-Tail communication."""
549 def attention_call(
550 query: Any, key: Any, value: Any,
551 call_kwargs: dict[str, Any]) -> Any:
552 """Invoke the original QKV forward with transformed tensors."""
553 return original_forward(query, key, value, **call_kwargs)
555 keep_kwargs, peer_kwargs = _prepare_head_tail_sdpa_kwargs(
556 kwargs, q, k, cp_mesh
557 )
558 return head_tail_load_balance_attention(
559 attention_call,
560 q,
561 k,
562 v,
563 keep_kwargs,
564 cp_mesh,
565 peer_attention_kwargs=peer_kwargs,
566 )
568 return cp_forward
571@inner_wrapper
572def sdpa_hf_load_balance_cp_wrapper(
573 target_module: Module, mesh: Any, tp_mesh: Any,
574 cp_mesh: Any, ep_mesh: Any) -> Callable:
575 """Apply Colossal Head-Tail load balancing at an HF SDPA primitive.
577 Args:
578 target_module: HF-style attention containing an SDPA primitive call.
579 mesh: Framework-owned root mesh; unused by this wrapper.
580 tp_mesh: Framework-owned TP mesh; unused by this wrapper.
581 cp_mesh: Framework-owned CP communication mesh.
582 ep_mesh: Framework-owned EP mesh; unused by this wrapper.
583 """
584 del mesh, tp_mesh, ep_mesh
585 if cp_mesh is None or cp_mesh.size() <= 1:
586 raise ValueError(
587 "sdpa_hf_load_balance_cp_wrapper requires an active CP mesh"
588 )
589 original_forward = target_module.forward
590 original_sdpa = F.scaled_dot_product_attention
592 @functools.wraps(original_forward)
593 def cp_forward(
594 hidden_states: Any, *args: Any, **kwargs: Any) -> Any:
595 """Intercept the module's SDPA primitive for one forward call."""
596 fired = {"hit": False}
598 def load_balanced_sdpa(
599 q: Any, k: Any, v: Any,
600 **call_kwargs: Any) -> Any:
601 """Run one intercepted SDPA call with Head-Tail communication."""
602 fired["hit"] = True
603 keep_kwargs, peer_kwargs = _prepare_head_tail_sdpa_kwargs(
604 call_kwargs, q, k, cp_mesh
605 )
607 def _call_original_sdpa(query, key, value, attention_kwargs):
608 return original_sdpa(query, key, value, **attention_kwargs)
610 return head_tail_load_balance_attention(
611 _call_original_sdpa,
612 q,
613 k,
614 v,
615 keep_kwargs,
616 cp_mesh,
617 peer_attention_kwargs=peer_kwargs,
618 )
620 F.scaled_dot_product_attention = load_balanced_sdpa
621 try:
622 output = original_forward(hidden_states, *args, **kwargs)
623 finally:
624 F.scaled_dot_product_attention = original_sdpa
625 if not fired["hit"]:
626 raise RuntimeError(
627 "CP wrapper 'sdpa_hf_load_balance' did not intercept any "
628 "F.scaled_dot_product_attention call on "
629 f"{type(target_module).__name__}; choose the QKV wrapper or "
630 "provide a matching custom inner_wrapper"
631 )
632 return output
634 return cp_forward
637# ────────────────────────────────────────────────────────────────────────────
638# Pure Ulysses attention wrappers (explicit plan_overrides methods)
639# ────────────────────────────────────────────────────────────────────────────
641@inner_wrapper
642def sdpa_qkv_ulysses_cp_wrapper(
643 target_module: Module, mesh: Any, tp_mesh: Any,
644 cp_mesh: Any, ep_mesh: Any) -> Callable:
645 """Apply Pure Ulysses to a separated ``forward(query, key, value, ...)``."""
646 del mesh, tp_mesh, ep_mesh
647 _require_ulysses_cp_mesh(cp_mesh, "sdpa_qkv_ulysses")
648 original_forward = target_module.forward
650 @functools.wraps(original_forward)
651 def cp_forward(*args: Any, **kwargs: Any) -> Any:
652 """Run the original QKV forward in the Ulysses head-sharded layout."""
653 return _ulysses_qkv_forward(original_forward, cp_mesh, args, kwargs)
655 return cp_forward
658@inner_wrapper
659def flex_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 FlexAttention QKV forward."""
663 del mesh, tp_mesh, ep_mesh
664 _require_ulysses_cp_mesh(cp_mesh, "flex_qkv_ulysses")
665 original_forward = target_module.forward
667 @functools.wraps(original_forward)
668 def cp_forward(*args: Any, **kwargs: Any) -> Any:
669 """Run the original FlexAttention QKV forward in the Ulysses layout."""
670 return _ulysses_qkv_forward(original_forward, cp_mesh, args, kwargs)
672 return cp_forward
675@inner_wrapper
676def sdpa_hf_ulysses_cp_wrapper(
677 target_module: Module, mesh: Any, tp_mesh: Any,
678 cp_mesh: Any, ep_mesh: Any) -> Callable:
679 """Apply Pure Ulysses by intercepting HF SDPA primitive calls."""
680 del mesh, tp_mesh, ep_mesh
681 _require_ulysses_cp_mesh(cp_mesh, "sdpa_hf_ulysses")
682 original_forward = target_module.forward
683 original_sdpa = F.scaled_dot_product_attention
685 @functools.wraps(original_forward)
686 def cp_forward(*args: Any, **kwargs: Any) -> Any:
687 """Run the HF forward with F.sdpa temporarily Ulysses-exchanged."""
688 fired = {"hit": False}
690 def ulysses_sdpa(
691 query: Any, key: Any, value: Any,
692 **attention_kwargs: Any) -> Any:
693 """Run one intercepted SDPA call in the Ulysses layout."""
694 fired["hit"] = True
695 return _ulysses_attention_call(
696 original_sdpa,
697 cp_mesh,
698 query,
699 key,
700 value,
701 attention_kwargs,
702 )
704 F.scaled_dot_product_attention = ulysses_sdpa
705 try:
706 output = original_forward(*args, **kwargs)
707 finally:
708 F.scaled_dot_product_attention = original_sdpa
709 if not fired["hit"]:
710 raise RuntimeError(
711 "CP wrapper 'sdpa_hf_ulysses' did not intercept any "
712 "F.scaled_dot_product_attention call; the selected wrapper "
713 "does not match the module implementation"
714 )
715 return output
717 return cp_forward
720@inner_wrapper
721def flex_hf_ulysses_cp_wrapper(
722 target_module: Module, mesh: Any, tp_mesh: Any,
723 cp_mesh: Any, ep_mesh: Any) -> Callable:
724 """Apply Pure Ulysses by intercepting HF FlexAttention primitive calls."""
725 del mesh, tp_mesh, ep_mesh
726 _require_ulysses_cp_mesh(cp_mesh, "flex_hf_ulysses")
727 original_forward = target_module.forward
728 flex_attention_module = importlib.import_module(
729 "torch.nn.attention.flex_attention"
730 )
731 original_flex_attention = flex_attention_module.flex_attention
733 @functools.wraps(original_forward)
734 def cp_forward(*args: Any, **kwargs: Any) -> Any:
735 """Run the HF forward with flex_attention temporarily Ulysses-exchanged."""
736 fired = {"hit": False}
738 def ulysses_flex(
739 query: Any, key: Any, value: Any,
740 **attention_kwargs: Any) -> Any:
741 """Run one intercepted FlexAttention call in the Ulysses layout."""
742 fired["hit"] = True
743 return _ulysses_attention_call(
744 original_flex_attention,
745 cp_mesh,
746 query,
747 key,
748 value,
749 attention_kwargs,
750 )
752 flex_attention_module.flex_attention = ulysses_flex
753 try:
754 output = original_forward(*args, **kwargs)
755 finally:
756 flex_attention_module.flex_attention = original_flex_attention
757 if not fired["hit"]:
758 raise RuntimeError(
759 "CP wrapper 'flex_hf_ulysses' did not intercept any "
760 "flex_attention call; the selected wrapper does not match "
761 "the module implementation"
762 )
763 return output
765 return cp_forward
768# ────────────────────────────────────────────────────────────────────────────
769# Hybrid attention wrappers (explicit Target with wrapper-local degree)
770# ────────────────────────────────────────────────────────────────────────────
772def _validate_hybrid_config(
773 cp_mesh: Any, ulysses_degree: int, wrapper_name: str) -> None:
774 """Validate Hybrid topology at wrapper installation time."""
775 if cp_mesh is None or cp_mesh.size() <= 1:
776 raise ValueError(
777 f"Hybrid wrapper {wrapper_name!r} requires an active CP mesh"
778 )
779 if isinstance(ulysses_degree, bool) or not isinstance(ulysses_degree, int):
780 raise TypeError(
781 f"Hybrid wrapper {wrapper_name!r} requires integer "
782 f"ulysses_degree, got {type(ulysses_degree).__name__}"
783 )
784 cp_size = cp_mesh.size()
785 if not 1 < ulysses_degree < cp_size:
786 raise ValueError(
787 f"Hybrid wrapper {wrapper_name!r} requires 1 < ulysses_degree "
788 f"< cp_size, got ulysses_degree={ulysses_degree}, cp_size={cp_size}"
789 )
790 if cp_size % ulysses_degree:
791 raise ValueError(
792 f"cp_size ({cp_size}) must be divisible by ulysses_degree "
793 f"({ulysses_degree})"
794 )
797def _apply_qkv_hybrid_wrapper(
798 target_module, cp_mesh, ulysses_degree, wrapper_name, *,
799 prepare_sdpa_mask: bool):
800 """Replace a separated QKV forward with local-tensor Hybrid CP."""
801 _validate_hybrid_config(cp_mesh, ulysses_degree, wrapper_name)
802 original_forward = target_module.forward
804 @functools.wraps(original_forward)
805 def cp_forward(q: Any, k: Any, v: Any, **kwargs: Any) -> Any:
806 """Run the original QKV forward through local-tensor Hybrid CP."""
807 def attention_call(
808 query: Any, key: Any, value: Any,
809 call_kwargs: dict[str, Any]) -> Any:
810 """Invoke the original QKV forward with transformed tensors."""
811 return original_forward(query, key, value, **call_kwargs)
813 if prepare_sdpa_mask:
814 kwargs = _prepare_hybrid_sdpa_kwargs(
815 kwargs, q, k, cp_mesh, ulysses_degree
816 )
817 return hybrid_cp_attention(
818 attention_call,
819 q,
820 k,
821 v,
822 kwargs,
823 cp_mesh,
824 ulysses_degree,
825 )
827 return cp_forward
830@inner_wrapper
831def sdpa_qkv_hybrid_cp_wrapper(
832 target_module: Module, mesh: Any, tp_mesh: Any, cp_mesh: Any,
833 ep_mesh: Any, ulysses_degree: int) -> Callable:
834 """Apply Hybrid CP to a separated SDPA ``forward(query, key, value, ...)``."""
835 del mesh, tp_mesh, ep_mesh
836 return _apply_qkv_hybrid_wrapper(
837 target_module,
838 cp_mesh,
839 ulysses_degree,
840 "sdpa_qkv_hybrid",
841 prepare_sdpa_mask=True,
842 )
845@inner_wrapper
846def flex_qkv_hybrid_cp_wrapper(
847 target_module: Module, mesh: Any, tp_mesh: Any, cp_mesh: Any,
848 ep_mesh: Any, ulysses_degree: int) -> Callable:
849 """Apply Hybrid CP to a separated FlexAttention QKV forward."""
850 del mesh, tp_mesh, ep_mesh
851 return _apply_qkv_hybrid_wrapper(
852 target_module,
853 cp_mesh,
854 ulysses_degree,
855 "flex_qkv_hybrid",
856 prepare_sdpa_mask=False,
857 )
860@inner_wrapper
861def sdpa_hf_hybrid_cp_wrapper(
862 target_module: Module, mesh: Any, tp_mesh: Any, cp_mesh: Any,
863 ep_mesh: Any, ulysses_degree: int) -> Callable:
864 """Apply Hybrid CP by intercepting HF SDPA primitive calls."""
865 del mesh, tp_mesh, ep_mesh
866 _validate_hybrid_config(cp_mesh, ulysses_degree, "sdpa_hf_hybrid")
867 original_forward = target_module.forward
868 original_sdpa = F.scaled_dot_product_attention
870 @functools.wraps(original_forward)
871 def cp_forward(*args: Any, **kwargs: Any) -> Any:
872 """Run the HF attention forward with temporary Hybrid SDPA interception."""
873 fired = {"hit": False}
875 def hybrid_sdpa(
876 query: Any, key: Any, value: Any,
877 **attention_kwargs: Any) -> Any:
878 """Route one intercepted SDPA call through Hybrid CP."""
879 fired["hit"] = True
880 attention_kwargs = _prepare_hybrid_sdpa_kwargs(
881 attention_kwargs,
882 query,
883 key,
884 cp_mesh,
885 ulysses_degree,
886 )
888 def _call_original_sdpa(call_query, call_key, call_value, call_kwargs):
889 return original_sdpa(
890 call_query, call_key, call_value, **call_kwargs
891 )
893 return hybrid_cp_attention(
894 _call_original_sdpa,
895 query,
896 key,
897 value,
898 attention_kwargs,
899 cp_mesh,
900 ulysses_degree,
901 )
903 F.scaled_dot_product_attention = hybrid_sdpa
904 try:
905 output = original_forward(*args, **kwargs)
906 finally:
907 F.scaled_dot_product_attention = original_sdpa
908 if not fired["hit"]:
909 raise RuntimeError(
910 "CP wrapper 'sdpa_hf_hybrid' did not intercept any "
911 "F.scaled_dot_product_attention call; the selected wrapper "
912 "does not match the module implementation"
913 )
914 return output
916 return cp_forward
919@inner_wrapper
920def flex_hf_hybrid_cp_wrapper(
921 target_module: Module, mesh: Any, tp_mesh: Any, cp_mesh: Any,
922 ep_mesh: Any, ulysses_degree: int) -> Callable:
923 """Apply Hybrid CP by intercepting HF FlexAttention primitive calls."""
924 del mesh, tp_mesh, ep_mesh
925 _validate_hybrid_config(cp_mesh, ulysses_degree, "flex_hf_hybrid")
926 original_forward = target_module.forward
927 flex_attention_module = importlib.import_module(
928 "torch.nn.attention.flex_attention"
929 )
930 original_flex_attention = flex_attention_module.flex_attention
932 @functools.wraps(original_forward)
933 def cp_forward(*args: Any, **kwargs: Any) -> Any:
934 """Run the HF attention forward with temporary Hybrid Flex interception."""
935 fired = {"hit": False}
937 def hybrid_flex(
938 query: Any, key: Any, value: Any,
939 **attention_kwargs: Any) -> Any:
940 """Route one intercepted FlexAttention call through Hybrid CP."""
941 fired["hit"] = True
943 def _call_original_flex(call_query, call_key, call_value, call_kwargs):
944 return original_flex_attention(
945 call_query, call_key, call_value, **call_kwargs
946 )
948 return hybrid_cp_attention(
949 _call_original_flex,
950 query,
951 key,
952 value,
953 attention_kwargs,
954 cp_mesh,
955 ulysses_degree,
956 )
958 flex_attention_module.flex_attention = hybrid_flex
959 try:
960 output = original_forward(*args, **kwargs)
961 finally:
962 flex_attention_module.flex_attention = original_flex_attention
963 if not fired["hit"]:
964 raise RuntimeError(
965 "CP wrapper 'flex_hf_hybrid' did not intercept any "
966 "flex_attention call; the selected wrapper does not match "
967 "the module implementation"
968 )
969 return output
971 return cp_forward
974# ────────────────────────────────────────────────────────────────────────────
975# MLA/DSA Ulysses wrapper
976# ────────────────────────────────────────────────────────────────────────────
978def _input_cp_sharding(text_model, context):
979 """Build the CP-rank input-slicing rewrite request for a text model."""
980 if getattr(text_model, _ULYSSES_WRAPPED_FLAG, False):
981 return None
982 original_forward = text_model.forward
984 @functools.wraps(original_forward)
985 def forward_with_sequence_sharding(*args: Any, **kwargs: Any) -> Any:
986 """Slice sequence-dim inputs to this CP rank before the forward."""
987 if getattr(text_model, "_hyper_cp_inside_forward", False):
988 return original_forward(*args, **kwargs)
989 call_kwargs = kwargs.copy()
990 inputs_embeds = call_kwargs.get("inputs_embeds")
991 if inputs_embeds is None:
992 return original_forward(*args, **kwargs)
993 call_kwargs["inputs_embeds"] = _slice_sequence(inputs_embeds, 1, context)
994 position_ids = call_kwargs.get("position_ids")
995 if position_ids is not None:
996 call_kwargs["position_ids"] = _slice_sequence(
997 position_ids, position_ids.ndim - 1, context)
998 mome_mask = call_kwargs.get("mome_mask")
999 if mome_mask is not None:
1000 call_kwargs["mome_mask"] = _slice_sequence(mome_mask, 1, context)
1001 if call_kwargs.get("use_cache"):
1002 raise ValueError("Ulysses CP requires use_cache=False")
1003 call_kwargs["use_cache"] = False
1004 setattr(text_model, "_hyper_cp_inside_forward", True)
1005 try:
1006 return original_forward(*args, **call_kwargs)
1007 finally:
1008 setattr(text_model, "_hyper_cp_inside_forward", False)
1010 return _ForwardRewriteRequest(
1011 text_model,
1012 forward_with_sequence_sharding,
1013 companion_attrs={_ULYSSES_WRAPPED_FLAG: True},
1014 )
1017def _validate_ulysses_requirements(target_module, cp_size):
1018 """Validate the model configuration supported by Ulysses CP."""
1019 config = getattr(target_module, "config", None)
1020 if config is None:
1021 raise ValueError(
1022 "MLA/DSA Ulysses wrapper requires target_module.config")
1023 text_config = getattr(config, "text_config", config)
1024 if text_config is None:
1025 raise ValueError(
1026 "MLA/DSA Ulysses wrapper requires a non-None text_config")
1027 for name in ("num_attention_heads", "index_num_attention_heads"):
1028 value = getattr(text_config, name, None)
1029 if value is None:
1030 raise ValueError(
1031 "MLA/DSA Ulysses wrapper requires "
1032 f"text_config.{name}")
1033 count = int(value)
1034 if count % cp_size:
1035 raise ValueError(
1036 f"{name}={count} is not divisible by CP size {cp_size}")
1037 if getattr(text_config, "dsa_dense_warm_up", False):
1038 raise ValueError("MLA/DSA CP does not support DSA dense warm-up")
1039 if not getattr(text_config, "apply_fa_rescale", False):
1040 raise ValueError("MLA/DSA CP requires apply_fa_rescale=True")
1041 if getattr(text_config, "use_fused_sink_fa", False):
1042 raise ValueError("MLA/DSA CP does not support fused sink FA")
1045def _collect_mla_dsa_text_models(target_module):
1046 """Collect text-model modules supported by the MLA/DSA wrapper.
1048 Args:
1049 target_module: Root module whose descendants are inspected.
1051 Returns:
1052 The matching text-model modules in traversal order.
1053 """
1054 text_models = []
1055 for name, module in target_module.named_modules():
1056 if name.rsplit(".", maxsplit=1)[-1] in {"text_model", "language_model"}:
1057 text_models.append(module)
1058 return text_models
1061def _collect_mla_dsa_attention_registries(attention_module):
1062 """Collect registries that expose both required MLA/DSA backends.
1064 Args:
1065 attention_module: Python module that owns the attention registries.
1067 Returns:
1068 Dictionaries containing both required attention functions.
1069 """
1070 attention_registries = []
1071 required_attention_functions = {"npu_fa_rescale", "dsa_sparse_attention"}
1072 for value in vars(attention_module).values():
1073 if isinstance(value, dict) and required_attention_functions <= value.keys():
1074 attention_registries.append(value)
1075 return attention_registries
1078@inner_wrapper
1079def mla_dsa_ulysses_cp_wrapper( # pylint: disable=inconsistent-return-statements
1080 target_module: Module, mesh: Any, tp_mesh: Any,
1081 cp_mesh: Any, ep_mesh: Any) -> Any:
1082 """Configure input, MoME, MLA and DSA Ulysses adaptations."""
1083 del mesh, tp_mesh, ep_mesh
1084 if cp_mesh is None or cp_mesh.size() <= 1:
1085 raise ValueError(
1086 "MLA/DSA Ulysses wrapper requires an active CP mesh")
1087 if getattr(target_module, _ULYSSES_WRAPPED_FLAG, False):
1088 return
1089 _validate_ulysses_requirements(target_module, cp_mesh.size())
1090 context = _UlyssesContext(cp_mesh)
1092 text_models = _collect_mla_dsa_text_models(target_module)
1093 if not text_models:
1094 raise RuntimeError("Cannot find a text or language model")
1095 requests = []
1096 for text_model in text_models:
1097 request = _input_cp_sharding(text_model, context)
1098 if request is not None:
1099 requests.append(request)
1101 attention_modules = {
1102 inspect.getmodule(module) for module in target_module.modules()
1103 if getattr(module, "attention_type", None) in {"mla", "dsa"}}
1104 attention_modules.discard(None)
1105 if len(attention_modules) != 1:
1106 names = sorted(module.__name__ for module in attention_modules)
1107 raise RuntimeError(
1108 f"Expected one MLA/DSA attention module, found {names}")
1109 attention_module = next(iter(attention_modules))
1110 attention_registries = _collect_mla_dsa_attention_registries(attention_module)
1111 if len(attention_registries) != 1:
1112 raise RuntimeError(
1113 "Expected one attention-function registry containing MLA and DSA backends")
1114 attention_functions = attention_registries[0]
1115 _mome_cp_halo_exchange(attention_module, context)
1116 _mla_cp_alltoall(attention_functions, context)
1117 _dsa_cp_alltoall(attention_module, attention_functions, context)
1119 original_forward = target_module.forward
1121 @functools.wraps(original_forward)
1122 def forward_with_ulysses_adapters(*args: Any, **kwargs: Any) -> Any:
1123 """Pass the forward through; the adapters are installed out of band."""
1124 return original_forward(*args, **kwargs)
1126 requests.append(_ForwardRewriteRequest(
1127 target_module,
1128 forward_with_ulysses_adapters,
1129 companion_attrs={
1130 "_hyper_ulysses_context": context,
1131 _ULYSSES_WRAPPED_FLAG: True,
1132 },
1133 ))
1134 return requests
1137# {registry_name: wrapper_fn} -- inner-wrapper named registry (05 §4.4.2).
1138# The mechanism is not CP-gated (declaration means application); the in-repo
1139# reference implementations carry CP semantics (their self-checks require an
1140# active cp axis); users may register arbitrary named schemes.
1141# Contract: @inner_wrapper fn(target_module, mesh, tp_mesh, cp_mesh,
1142# ep_mesh) RETURNS the replacement forward (K/V all-gather + dual-mode
1143# tolerance); the forward rewriter installs it (external wrappers may
1144# still replace in place and return None). Users may register their own
1145# named schemes:
1146# INNER_WRAPPER_REGISTRY["my_flash"] = my_fn, after which
1147# spec.inner_wrapper="my_flash" references it by name.
1148INNER_WRAPPER_REGISTRY = {
1149 "sdpa_qkv": sdpa_qkv_cp_wrapper, # NeMo convention forward(q,k,v,...) + D-04 mask
1150 "sdpa_hf": sdpa_hf_cp_wrapper, # HF convention + F.sdpa primitive interception (misfire detection)
1151 "flex_qkv": flex_qkv_cp_wrapper,
1152 "flex_hf": flex_hf_cp_wrapper,
1153 "sdpa_qkv_ulysses": sdpa_qkv_ulysses_cp_wrapper,
1154 "flex_qkv_ulysses": flex_qkv_ulysses_cp_wrapper,
1155 "sdpa_hf_ulysses": sdpa_hf_ulysses_cp_wrapper,
1156 "flex_hf_ulysses": flex_hf_ulysses_cp_wrapper,
1157 "mla_dsa_ulysses": mla_dsa_ulysses_cp_wrapper,
1158}
1160# Static requirements for shipped wrappers. Custom registry entries own their
1161# semantics; built-ins are known to contain CP collectives and therefore must
1162# run as black-box local regions during placement validation.
1163INNER_WRAPPER_REQUIREMENTS = {
1164 name: {
1165 "requires_cp": True,
1166 "region_dispatch": False,
1167 "forward_style": (
1168 "hf_hidden_states"
1169 if name.endswith("_hf") or "_hf_" in name
1170 else "qkv"
1171 ),
1172 }
1173 for name in INNER_WRAPPER_REGISTRY
1174}