« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Context-parallel scheduling for NoPE MLA with a required sigmoid gate."""
17from __future__ import annotations
19from typing import Any, Mapping
21# Model-owned gate math and CP scheduling use native Torch.
22import torch # pylint: disable=forbidden-backend-import
23from torch import nn # pylint: disable=forbidden-backend-import
25from hyper_parallel.distributed.context_parallel.collectives import async_ulysses_seq_to_head_launch
26from hyper_parallel.distributed.context_parallel.gated_mla_backward_overlap import (
27 GatedMLAInputState, GatedMLAOutputProjectionState,
28)
29from hyper_parallel.distributed.context_parallel.mla_context_parallel import MLAContextParallel
32def gated_mla_forward(module: nn.Module, runtime: MLAContextParallel, hidden_states: torch.Tensor,
33 attention_mask: torch.Tensor | None, options: Mapping[str, Any], *,
34 overlap_output_backward: bool = False) -> torch.Tensor:
35 """Run shared MLA inputs and keep the gate at its chosen overlap boundary.
37 Args:
38 module: Gated MLA owning all projection parameters.
39 runtime: Shared MLA CP planner and attention execution.
40 hidden_states: Chronological local tokens in BSH layout.
41 attention_mask: Optional mask supported by the runtime.
42 options: Packed lengths, positions and attention options.
43 overlap_output_backward: Cover dO A2A with output wgrad. Input backward
44 additionally requires pure Ulysses; hybrid inputs retain native autograd.
46 Returns:
47 Gated and projected hidden states in the original local-token order.
48 """
49 # NoPE does not consume caller frequencies, including their layout metadata.
50 # The module-owned identity callback also governs both CP projection paths.
51 options = {name: value for name, value in options.items() if name != "position_embeddings"}
52 # Input Q/K/V pairing is implemented only for pure Ulysses. Hybrid
53 # paths can still overlap output dO A2A independently of their input adjoints.
54 overlap_input_backward = overlap_output_backward and runtime.plan.all_gather_degree == 1
55 input_gate = _prepare_input_gate(module, runtime, hidden_states, overlap_input_backward)
56 if input_gate is not None:
57 overlap_input_backward = input_gate.defer_gate_gradient
58 attention_output, _ = runtime(
59 module, *module.project_latent_inputs(hidden_states), attention_mask=attention_mask,
60 restore_output=False, overlap_input_backward=overlap_input_backward, input_gate=input_gate, **options,
61 )
62 output_a2a_handle = None
63 if input_gate is None and runtime.plan.overlap and runtime.plan.ulysses_degree > 1:
64 # Launch before the native gate so its autograd graph and FSDP hooks
65 # stay outside the output Function while O communication is in flight.
66 output_a2a_handle = async_ulysses_seq_to_head_launch(attention_output, 2, 1, runtime.ulysses_mesh)
67 if input_gate is None:
68 gate = module.g_proj(hidden_states).sigmoid()
69 else:
70 gate, input_gate.value = input_gate.value, None
71 if overlap_output_backward:
72 deferred_gate = input_gate if input_gate is not None and input_gate.defer_gate_gradient else None
73 state = GatedMLAOutputProjectionState(
74 runtime.ulysses_mesh, module.v_head_dim, input_gate=deferred_gate, forward_a2a_handle=output_a2a_handle,
75 )
76 return module.o_proj(attention_output, mla_output_state=state, mla_output_gate=gate)
77 output = output_a2a_handle.wait() if output_a2a_handle is not None else runtime.restore_output(attention_output)
78 output = output.flatten(2)
79 return module.o_proj(output * gate)
82def _prepare_input_gate(module, runtime, hidden_states, overlap_input_backward):
83 """Use input-side gating only for the expanded pure-Ulysses schedule."""
84 plan = runtime.plan
85 if not plan.overlap or plan.ulysses_degree <= 1:
86 return None
87 if plan.strategy != "expanded_ulysses" or plan.all_gather_degree != 1:
88 return None
89 # Gate-gradient handoff needs the input dependency. Frozen inputs use
90 # native gate backward; inference does not build a backward graph.
91 defer_gate_gradient = overlap_input_backward and hidden_states.requires_grad
92 return GatedMLAInputState(
93 hidden_states, lambda source: module.g_proj(source).sigmoid(), defer_gate_gradient,
94 )