« 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"""Gate-specific input dependencies and output derivatives for MLA overlap."""
17from __future__ import annotations
19from collections.abc import Callable
20from dataclasses import dataclass
21from functools import wraps
22from typing import Any
24# Gate math and its autograd scheduling use native Torch.
25import torch # pylint: disable=forbidden-backend-import
26from torch import nn # pylint: disable=forbidden-backend-import
27from torch.autograd.function import once_differentiable # pylint: disable=forbidden-backend-import
28from torch.nn import functional as F # pylint: disable=forbidden-backend-import
30from hyper_parallel.distributed._builder.forward_rewriter import _ForwardRewriteRequest
31from hyper_parallel.distributed.context_parallel.collectives import (
32 AsyncCPCollective, async_ulysses_seq_to_head_launch,
33)
34from hyper_parallel.distributed.context_parallel.mla_backward_overlap import (
35 MLAExpandedInputA2AState, _InputA2ALaunch, launch_input_a2a,
36)
39@dataclass
40class GatedMLAInputState:
41 """Invocation-local gate values and the optional deferred native derivative.
43 The adapter supplies the model-owned gate computation. It is released after
44 forward; the remaining state belongs only to this invocation's paired nodes.
45 With defer_gate_gradient, the output saves dG and the input consumer returns
46 it only after submitting the input-gradient A2As.
47 """
49 hidden: torch.Tensor | None
50 projection: Callable[[torch.Tensor], torch.Tensor] | None
51 defer_gate_gradient: bool
52 value: torch.Tensor | None = None
53 gradient: torch.Tensor | None = None
55 def project(self, source: torch.Tensor) -> torch.Tensor:
56 """Run the managed gate and release its temporary forward-only references.
58 Args:
59 source: Local hidden states, optionally carrying the input launch edge.
61 Returns:
62 The model's post-activation gate in local-token layout.
63 """
64 if self.projection is None:
65 raise RuntimeError("MLA input gate must be projected exactly once")
66 self.value = self.projection(source)
67 self.hidden, self.projection = None, None
68 return self.value
70 def take_gradient(self) -> torch.Tensor:
71 """Consume the output multiplication's deferred gate gradient once."""
72 if self.gradient is None:
73 raise RuntimeError("MLA output projection did not release its gate gradient")
74 gradient, self.gradient = self.gradient, None
75 return gradient
78class _GatedExpandedInputA2AWait(torch.autograd.Function):
79 """Release gate backward only after the attention input gradients are sent."""
81 @staticmethod
82 def forward(ctx: Any, gate: torch.Tensor, state: MLAExpandedInputA2AState,
83 gate_state: GatedMLAInputState) -> tuple:
84 """Consume Q/K/V after native gate compute, keeping its gradient edge.
86 Args:
87 ctx: Autograd context for communication and gate-gradient ownership.
88 gate: Post-sigmoid local-token gate, whose backward is deferred here.
89 state: Q/K/V transfers already submitted before gate compute.
90 gate_state: Invocation-local handoff from the gated output derivative.
92 Returns:
93 Completed attention inputs in query, key, value order.
94 """
95 # Output projection consumes the gate value. This argument keeps the
96 # edge on which this consumer will return the deferred dG.
97 del gate
98 ctx.state, ctx.gate_state = state, gate_state
99 return state.wait()
101 @staticmethod
102 @once_differentiable
103 def backward(ctx: Any, grad_query: torch.Tensor, grad_key: torch.Tensor,
104 grad_value: torch.Tensor) -> tuple:
105 """Launch dQ/dK/dV before delivering dG to native sigmoid/g_proj backward.
107 Args:
108 ctx: This invocation's communication and gate-gradient states.
109 grad_query: Head-local query gradient from attention.
110 grad_key: Head-local key gradient from attention.
111 grad_value: Head-local value gradient from attention.
113 Returns:
114 Gate gradient and no gradients for either state object.
115 """
116 ctx.state.launch_backward(grad_query, grad_key, grad_value)
117 return ctx.gate_state.take_gradient(), None, None
120def finish_gated_expanded_inputs(
121 query: torch.Tensor, state: MLAExpandedInputA2AState, gate_state: GatedMLAInputState,
122) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
123 """Cover the last Q A2A with the model-owned native gate projection.
125 Args:
126 query: Final token-local query, after V and K have been submitted.
127 state: This invocation's Q/K/V communication state.
128 gate_state: Model gate computation and its selected gradient handoff.
130 Returns:
131 Completed attention inputs in head-local layout; gate is saved in gate_state.
132 """
133 launch_input = _InputA2ALaunch.apply if gate_state.defer_gate_gradient else launch_input_a2a
134 gate_input = launch_input(query, gate_state.hidden, state.query)
135 gate = gate_state.project(gate_input)
136 if gate_state.defer_gate_gradient:
137 return _GatedExpandedInputA2AWait.apply(gate, state, gate_state)
138 return state.wait()
141@dataclass
142class GatedMLAOutputProjectionState:
143 """Own gated output communication and optional input-side gate gradients.
145 A pre-launched O transfer covers the native gate projection. For input-side
146 gating, input_gate instead delivers dG to the paired input consumer, which
147 releases the native gate backward after submitting the input gradients.
148 """
150 mesh: Any
151 head_dim: int
152 input_gate: GatedMLAInputState | None = None
153 forward_a2a_handle: AsyncCPCollective | None = None
156class _GatedOutputProjection(torch.autograd.Function):
157 """Own gated output multiplication, Linear and the inverse output A2A."""
159 @staticmethod
160 def forward(ctx: Any, attention_output: torch.Tensor, weight: torch.Tensor,
161 bias: torch.Tensor | None, gate: torch.Tensor,
162 state: GatedMLAOutputProjectionState) -> torch.Tensor:
163 """Restore the FA output, multiply its required gate and apply Wo.
165 Args:
166 ctx: Autograd context for operands, layout and autocast settings.
167 attention_output: FA output in [batch, sequence, local_heads, value_dim].
168 weight: Managed output-projection weight.
169 bias: Optional output-projection bias.
170 gate: Native post-sigmoid gate in token-local layout.
171 state: Output layout, O transfer and deferred gate consumer.
173 Returns:
174 Gated and projected local hidden states.
175 """
176 output_a2a_handle, state.forward_a2a_handle = state.forward_a2a_handle, None
177 if output_a2a_handle is None:
178 output_a2a_handle = async_ulysses_seq_to_head_launch(attention_output, 2, 1, state.mesh)
179 # This Function owns the collective's derivative. Native g_proj and
180 # sigmoid stay outside it so their module hooks and autograd still run.
181 local_output = output_a2a_handle.wait().flatten(2)
182 projected = local_output * gate
183 ctx.save_for_backward(local_output, projected, weight, gate)
184 ctx.state, ctx.has_bias = state, bias is not None
185 ctx.device_type = local_output.device.type
186 ctx.autocast_enabled = torch.is_autocast_enabled(ctx.device_type)
187 ctx.autocast_dtype = torch.get_autocast_dtype(ctx.device_type) if ctx.autocast_enabled else None
188 return F.linear(projected, weight, bias) # pylint: disable=not-callable
190 @staticmethod
191 @once_differentiable
192 def backward(ctx: Any, gradient: torch.Tensor) -> tuple[torch.Tensor | None, ...]:
193 """Overlap dO A2A with gate multiplication gradients and Wo wgrad.
195 Args:
196 ctx: Saved operands, autocast settings and output layout.
197 gradient: Gradient of the projected hidden states.
199 Returns:
200 FA output, weight, bias, gate and no state gradients. Input-side
201 gating delivers dG through its paired consumer instead.
202 """
203 with torch.autocast(ctx.device_type, enabled=ctx.autocast_enabled, dtype=ctx.autocast_dtype):
204 return _gated_output_projection_backward(ctx, gradient)
207def _gated_output_projection_backward(ctx, gradient):
208 """Keep the dO launch before wgrad and the wait before returning dF."""
209 local_output, projected, weight, gate = ctx.saved_tensors
210 state = ctx.state
211 needs_attention_grad, needs_weight_grad, needs_bias_grad, needs_gate_grad = ctx.needs_input_grad[:4]
212 needs_gate = needs_gate_grad or state.input_gate is not None
213 grad_projected = gradient.matmul(weight).to(projected.dtype) if needs_attention_grad or needs_gate else None
214 grad_output_a2a = None
215 if needs_attention_grad:
216 grad_local_output = (grad_projected * gate).reshape(*local_output.shape[:2], -1, state.head_dim)
217 grad_output_a2a = async_ulysses_seq_to_head_launch(grad_local_output, 1, 2, state.mesh)
218 grad_gate = grad_projected * local_output if needs_gate else None
219 if state.input_gate is not None:
220 state.input_gate.gradient = grad_gate
221 grad_gate = None
222 flat_input = projected.reshape(-1, projected.shape[-1])
223 flat_gradient = gradient.reshape(-1, gradient.shape[-1])
224 grad_weight = None
225 if needs_weight_grad:
226 grad_weight = flat_input.transpose(0, 1).matmul(flat_gradient).transpose(0, 1)
227 grad_bias = flat_gradient.sum(0) if ctx.has_bias and needs_bias_grad else None
228 grad_attention_output = grad_output_a2a.wait() if grad_output_a2a is not None else None
229 return grad_attention_output, grad_weight, grad_bias, grad_gate, None
232def gated_mla_output_projection_request(projection: nn.Linear) -> _ForwardRewriteRequest:
233 """Install the gated output derivative after managed parameter pre-hooks.
235 Args:
236 projection: Plain, unwrapped output Linear, before FSDP installation.
238 Returns:
239 An atomic child rewrite; ordinary Linear calls remain unchanged.
240 """
241 if type(projection) is not nn.Linear: # pylint: disable=unidiomatic-typecheck
242 raise ValueError("Gated MLA backward overlap requires a plain output Linear")
243 if getattr(projection.forward, "__func__", None) is not nn.Linear.forward:
244 raise ValueError("Gated MLA backward overlap requires an unwrapped plain output Linear")
245 original = projection.forward
247 @wraps(original)
248 def gated_output_projection(*args: torch.Tensor,
249 mla_output_state: GatedMLAOutputProjectionState | None = None,
250 mla_output_gate: torch.Tensor | None = None, **kwargs: Any) -> torch.Tensor:
251 """Select the scheduled gated boundary only for explicit internal calls.
253 Args:
254 *args: Positional arguments of the original Linear.
255 mla_output_state: Optional state selecting the scheduled output path.
256 mla_output_gate: Required post-sigmoid gate for a scheduled call.
257 **kwargs: Keyword arguments of the original Linear.
259 Returns:
260 Projected hidden states using ordinary or gated scheduled autograd.
261 """
262 if mla_output_state is None:
263 if mla_output_gate is not None:
264 raise ValueError("Scheduled gated MLA output requires an output state")
265 return original(*args, **kwargs)
266 if len(args) != 1 or kwargs:
267 raise TypeError("Scheduled gated MLA output projection requires one positional activation")
268 if mla_output_gate is None:
269 raise ValueError("Scheduled gated MLA output requires a gate tensor")
270 return _GatedOutputProjection.apply(
271 args[0], projection.weight, projection.bias, mla_output_gate, mla_output_state,
272 )
274 return _ForwardRewriteRequest(projection, gated_output_projection)