Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / attention.py: 52%
145 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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# ============================================================================
16"""context_parallel.attention: CP attention-layout adaptations.
18Model-side CP adaptations on top of the collectives in
19context_parallel/collectives.py: MoME halo exchange, MLA/DSA all-to-all
20layouts, head/tail load-balanced attention and the D-04 offset-aware
21causal mask.
23Split out of components/distributed/cp_utils.py in stage 4e.
24"""
26import contextvars
27import functools
28import inspect
29from dataclasses import dataclass
30from typing import Any, Callable, Optional
31import torch
32import torch.distributed as dist
33from torch import Tensor
34from torch.distributed.nn.functional import all_gather as differentiable_all_gather
35from hyper_parallel.core.utils import communication
36from hyper_parallel.distributed.context_parallel.collectives import (
37 _ULYSSES_WRAPPED_FLAG,
38 _gather_sequence,
39 _global_seq_len,
40 _head_to_sequence,
41 _sequence_to_head,
42 _slice_sink,
43 flex_cp_allgather,
44)
47@dataclass
48class _DSATensorContext:
49 query: Tensor
50 key: Tensor
51 q_pe: Tensor
52 k_pe: Tensor
55_dsa_tensor_context = contextvars.ContextVar(
56 "hyper_dsa_tensor_context", default=None)
58_DSA_KL_ARGUMENTS = (
59 "index_query", "index_key", "merge_weight", "query", "key",
60 "topk_indices", "softmax_max", "softmax_sum", "query_rope",
61 "key_rope", "actual_seq_qlen", "actual_seq_klen", "scale",
62 "loss_coeff",
63)
64_DSA_KL_SIGNATURE = inspect.Signature(
65 inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD)
66 for name in _DSA_KL_ARGUMENTS
67)
70def _mome_cp_halo_exchange(attention_module, context):
71 """Configure cross-rank halo exchange for MoME convolution."""
72 original = getattr(attention_module, "_apply_mome")
73 if getattr(original, _ULYSSES_WRAPPED_FLAG, False):
74 return
76 @functools.wraps(original)
77 def apply_mome_with_halo(
78 hidden_states: Tensor,
79 mome_mask: Tensor,
80 conv: Any,
81 use_fused: bool,
82 ) -> Tensor:
83 """Run MoME convolution with the previous rank's halo rows prepended."""
84 halo = conv.kernel_size[0] - 1
85 if halo == 0:
86 return original(hidden_states, mome_mask, conv, use_fused)
87 if hidden_states.size(1) < halo:
88 raise ValueError(
89 f"local sequence {hidden_states.size(1)} is shorter than "
90 f"MOME halo {halo}")
91 tails = differentiable_all_gather(
92 hidden_states[:, -halo:].contiguous(), group=context.group)
93 mask_tail = mome_mask[:, -halo:].to(
94 hidden_states.dtype).contiguous()
95 masks = differentiable_all_gather(mask_tail, group=context.group)
96 if context.rank == 0:
97 left_states = tails[-1] * 0
98 left_mask = torch.zeros_like(mask_tail, dtype=torch.bool)
99 else:
100 left_states = tails[context.rank - 1]
101 left_mask = masks[context.rank - 1].bool()
102 output = original(
103 torch.cat((left_states, hidden_states), dim=1),
104 torch.cat((left_mask, mome_mask.bool()), dim=1), conv, use_fused)
105 return output[:, halo:].contiguous()
107 setattr(apply_mome_with_halo, _ULYSSES_WRAPPED_FLAG, True)
108 setattr(attention_module, "_apply_mome", apply_mome_with_halo)
111def _mla_cp_alltoall(attention_functions, context):
112 """Configure CP all-to-all around the MLA backend."""
113 original = attention_functions["npu_fa_rescale"]
114 if getattr(original, _ULYSSES_WRAPPED_FLAG, False):
115 return
117 @functools.wraps(original)
118 def mla_with_sequence_head_exchange(
119 module: Any, query: Tensor, key: Tensor, value: Tensor,
120 attention_mask: Any, **kwargs: Any) -> Tensor:
121 """Run the MLA backend with Q/K/V exchanged to head-sharded layout."""
122 if module.attention_type != "mla":
123 return original(
124 module, query, key, value, attention_mask, **kwargs)
125 if not module.apply_fa_rescale or module.use_fused_sink_fa:
126 raise ValueError(
127 "MLA CP supports only non-fused npu_fa_rescale")
128 local_shape = tuple(query.shape)
129 query = _sequence_to_head(query, context)
130 key = _sequence_to_head(key, context)
131 value = _sequence_to_head(value, context)
132 length = query.size(1)
133 call_kwargs = kwargs.copy()
134 call_kwargs.update(
135 seq_length=length,
136 n_head=query.size(2),
137 actual_q_len=_global_seq_len(
138 kwargs.get("actual_q_len"), length, query.device),
139 actual_kv_len=_global_seq_len(
140 kwargs.get("actual_kv_len"), length, query.device),
141 param_sink_key=_slice_sink(
142 kwargs.get("param_sink_key"), context),
143 param_sink_value=_slice_sink(
144 kwargs.get("param_sink_value"), context),
145 )
146 output = _head_to_sequence(
147 original(
148 module, query, key, value, attention_mask, **call_kwargs),
149 context)
150 if output.shape[:3] != torch.Size(local_shape[:3]):
151 raise RuntimeError(
152 f"MLA CP output {tuple(output.shape)} does not restore "
153 f"{local_shape}")
154 return output
156 setattr(mla_with_sequence_head_exchange, _ULYSSES_WRAPPED_FLAG, True)
157 attention_functions["npu_fa_rescale"] = mla_with_sequence_head_exchange
160def _gather_dsa_teacher_heads(tensor: Tensor, head_dim: int, context: Any) -> Tensor:
161 """Restore all teacher heads without adding an auxiliary gradient path."""
162 return torch.cat(
163 differentiable_all_gather(tensor.detach().contiguous(), group=context.group), dim=head_dim)
166def _dsa_cp_alltoall(attention_module, attention_functions, context):
167 """Configure CP all-to-all for the DSA indexer, attention, and KL loss."""
168 original_indexer = attention_module.dsa_lightning_indexer_forward
169 original_sparse = attention_functions["dsa_sparse_attention"]
170 original_kl = attention_module.SparseLightningIndexerKLLossTrainFunction
172 if not getattr(original_indexer, _ULYSSES_WRAPPED_FLAG, False):
173 @functools.wraps(original_indexer)
174 def index_with_gathered_sequence(
175 module: Any, index_query: Tensor, index_key: Tensor,
176 merge_weight: Tensor,
177 actual_q_len: Any, actual_kv_len: Any) -> Any:
178 """Keep the full head reduction before the non-linear top-k."""
179 index_query = _gather_sequence(index_query, context)
180 merge_weight = _gather_sequence(merge_weight, context)
181 index_key = _gather_sequence(index_key, context)
182 length = index_query.size(1)
183 return original_indexer(
184 module, index_query, index_key, merge_weight,
185 _global_seq_len(
186 actual_q_len, length, index_query.device),
187 _global_seq_len(
188 actual_kv_len, length, index_query.device))
190 setattr(index_with_gathered_sequence, _ULYSSES_WRAPPED_FLAG, True)
191 attention_module.dsa_lightning_indexer_forward = (
192 index_with_gathered_sequence)
194 if not getattr(original_sparse, _ULYSSES_WRAPPED_FLAG, False):
195 @functools.wraps(original_sparse)
196 def sparse_attention_with_gathered_kv(
197 module: Any, query: Tensor, key: Tensor, value: Tensor,
198 attention_mask: Any, **kwargs: Any) -> tuple[Tensor, Any, Any]:
199 """Run DSA sparse attention with head-sharded Q and gathered K/V."""
200 del attention_mask
201 local_shape = tuple(query.shape)
202 query = _sequence_to_head(query, context)
203 q_pe = _sequence_to_head(kwargs["q_pe"], context)
204 key = _gather_sequence(key, context)
205 value = _gather_sequence(value, context)
206 k_pe = _gather_sequence(kwargs["k_pe"], context)
207 length = query.size(1)
208 call_kwargs = kwargs.copy()
209 call_kwargs.update(
210 q_pe=q_pe, k_pe=k_pe, seq_length=length,
211 n_head=query.size(2),
212 actual_q_len=_global_seq_len(
213 kwargs.get("actual_q_len"), length, query.device),
214 actual_kv_len=_global_seq_len(
215 kwargs.get("actual_kv_len"), length, query.device),
216 )
217 old_heads = module.num_heads
218 module.num_heads = query.size(2)
219 try:
220 output, softmax_max, softmax_sum = original_sparse(
221 module, query, key, value, None, **call_kwargs)
222 finally:
223 module.num_heads = old_heads
224 if module.training and not module.freeze_dsa:
225 _dsa_tensor_context.set(_DSATensorContext(
226 query=query, key=key, q_pe=q_pe, k_pe=k_pe))
227 output = _head_to_sequence(output, context)
228 if output.shape[:3] != torch.Size(local_shape[:3]):
229 raise RuntimeError(
230 f"DSA CP output {tuple(output.shape)} does not restore "
231 f"{local_shape}")
232 return output, softmax_max, softmax_sum
234 setattr(sparse_attention_with_gathered_kv, _ULYSSES_WRAPPED_FLAG, True)
235 attention_functions["dsa_sparse_attention"] = (
236 sparse_attention_with_gathered_kv)
238 if not getattr(original_kl, _ULYSSES_WRAPPED_FLAG, False):
239 class CPDSAKLLoss:
240 """Proxy the DSA KL loss with CP-transformed attention inputs."""
242 @staticmethod
243 def apply(*args: Any, **kwargs: Any) -> Any:
244 """Apply the original KL loss with saved global sequence tensors."""
245 arguments = _DSA_KL_SIGNATURE.bind(*args, **kwargs).arguments
246 saved = _dsa_tensor_context.get()
247 if saved is not None:
248 _dsa_tensor_context.set(None)
249 length = saved.query.size(1)
250 arguments.update(
251 query=_gather_dsa_teacher_heads(saved.query, 2, context).flatten(0, 1),
252 key=saved.key.flatten(0, 1),
253 query_rope=_gather_dsa_teacher_heads(saved.q_pe, 2, context).flatten(0, 1),
254 key_rope=saved.k_pe.flatten(0, 1),
255 softmax_max=_gather_dsa_teacher_heads(arguments["softmax_max"], -1, context),
256 softmax_sum=_gather_dsa_teacher_heads(arguments["softmax_sum"], -1, context),
257 actual_seq_qlen=_global_seq_len(
258 arguments["actual_seq_qlen"], length, saved.query.device),
259 actual_seq_klen=_global_seq_len(
260 arguments["actual_seq_klen"], length, saved.query.device),
261 )
262 return original_kl.apply(
263 *(arguments[name] for name in _DSA_KL_ARGUMENTS))
265 setattr(CPDSAKLLoss, _ULYSSES_WRAPPED_FLAG, True)
266 attention_module.SparseLightningIndexerKLLossTrainFunction = (
267 CPDSAKLLoss)
270def _head_tail_peer_rank(cp_mesh: Any) -> int:
271 """Return the mirror rank in the CP mesh."""
272 rank_list = list(cp_mesh.rank_list)
273 return rank_list[cp_mesh.size() - 1 - rank_list.index(dist.get_rank())]
276def _run_head_tail_half(
277 attention_fn: Callable[[Tensor, Tensor, Tensor, dict[str, Any]], Any],
278 query: Tensor, key: Tensor, value: Tensor,
279 call_kwargs: dict[str, Any]) -> Tensor:
280 """Run attention for one query half and check its output contract."""
281 output = attention_fn(query, key, value, call_kwargs)
282 if not isinstance(output, Tensor):
283 raise TypeError(
284 "Head-Tail load balance requires the attention callable to "
285 f"return a Tensor, got {type(output).__name__}"
286 )
287 return output
290def head_tail_load_balance_attention(
291 attention_fn: Callable[[Tensor, Tensor, Tensor, dict[str, Any]], Any],
292 query: Tensor, key: Tensor, value: Tensor,
293 attention_kwargs: dict[str, Any], cp_mesh: Any, *,
294 peer_attention_kwargs: Optional[dict[str, Any]] = None) -> Any:
295 """Run local-tensor Colossal Head-Tail communication.
297 The caller prepares split-specific mask or position metadata. When
298 ``peer_attention_kwargs`` is omitted, both attention calls reuse
299 ``attention_kwargs`` for backward compatibility.
300 """
301 if cp_mesh is None or cp_mesh.size() <= 1:
302 raise ValueError("Head-Tail load balance requires an active CP mesh")
303 local_q_len = query.shape[2]
304 if local_q_len % 2:
305 raise ValueError(
306 "Head-Tail load balance requires an even local Q sequence "
307 f"length, got {local_q_len}; pad the global sequence to a "
308 f"multiple of 2 * cp_size ({2 * cp_mesh.size()})"
309 )
311 peer_rank = _head_tail_peer_rank(cp_mesh)
312 query_peer = communication.p2p_exchange(
313 query.narrow(2, local_q_len // 2, local_q_len // 2), peer_rank)
314 global_key, global_value = flex_cp_allgather(key, value, 2, cp_mesh)
315 keep_output = _run_head_tail_half(
316 attention_fn, query.narrow(2, 0, local_q_len // 2),
317 global_key, global_value, attention_kwargs)
318 peer_output = _run_head_tail_half(
319 attention_fn, query_peer, global_key, global_value,
320 attention_kwargs if peer_attention_kwargs is None else peer_attention_kwargs,
321 )
322 return torch.cat(
323 [keep_output, communication.p2p_exchange(peer_output, peer_rank)], dim=2)
326def _cp_offset_causal_mask(q_len: int, kv_len: int, lo: int,
327 device, dtype=torch.bool):
328 """D-04: offset-aware causal mask (this rank's Q chunk has global offset lo).
330 Attendable positions: j <= lo + i (i is the local Q row index).
331 Replaces is_causal=True -- torch SDPA's is_causal is top-left aligned when
332 q_len != kv_len (equivalent to assuming Q starts at global position 0), so
333 under CP the chunks of rank>0 would be incorrectly masked (G4).
334 """
335 i = torch.arange(q_len, device=device).view(-1, 1)
336 j = torch.arange(kv_len, device=device).view(1, -1)
337 return (j <= (lo + i)).to(dtype)