Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / attention.py: 43%
143 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# ============================================================================
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 _dsa_cp_alltoall(attention_module, attention_functions, context):
161 """Configure CP all-to-all for the DSA indexer, attention, and KL loss."""
162 original_indexer = attention_module.dsa_lightning_indexer_forward
163 original_sparse = attention_functions["dsa_sparse_attention"]
164 original_kl = attention_module.SparseLightningIndexerKLLossTrainFunction
166 if not getattr(original_indexer, _ULYSSES_WRAPPED_FLAG, False):
167 @functools.wraps(original_indexer)
168 def index_with_gathered_sequence(
169 module: Any, index_query: Tensor, index_key: Tensor,
170 merge_weight: Tensor,
171 actual_q_len: Any, actual_kv_len: Any) -> Any:
172 """Run the DSA indexer on the CP-exchanged sequence layout."""
173 index_query = _sequence_to_head(index_query, context)
174 merge_weight = _sequence_to_head(
175 merge_weight.unsqueeze(-1), context).squeeze(-1)
176 index_key = _gather_sequence(index_key, context)
177 length = index_query.size(1)
178 return original_indexer(
179 module, index_query, index_key, merge_weight,
180 _global_seq_len(
181 actual_q_len, length, index_query.device),
182 _global_seq_len(
183 actual_kv_len, length, index_query.device))
185 setattr(index_with_gathered_sequence, _ULYSSES_WRAPPED_FLAG, True)
186 attention_module.dsa_lightning_indexer_forward = (
187 index_with_gathered_sequence)
189 if not getattr(original_sparse, _ULYSSES_WRAPPED_FLAG, False):
190 @functools.wraps(original_sparse)
191 def sparse_attention_with_gathered_kv(
192 module: Any, query: Tensor, key: Tensor, value: Tensor,
193 attention_mask: Any, **kwargs: Any) -> tuple[Tensor, Any, Any]:
194 """Run DSA sparse attention with head-sharded Q and gathered K/V."""
195 del attention_mask
196 local_shape = tuple(query.shape)
197 query = _sequence_to_head(query, context)
198 q_pe = _sequence_to_head(kwargs["q_pe"], context)
199 key = _gather_sequence(key, context)
200 value = _gather_sequence(value, context)
201 k_pe = _gather_sequence(kwargs["k_pe"], context)
202 length = query.size(1)
203 call_kwargs = kwargs.copy()
204 call_kwargs.update(
205 q_pe=q_pe, k_pe=k_pe, seq_length=length,
206 n_head=query.size(2),
207 actual_q_len=_global_seq_len(
208 kwargs.get("actual_q_len"), length, query.device),
209 actual_kv_len=_global_seq_len(
210 kwargs.get("actual_kv_len"), length, query.device),
211 )
212 old_heads = module.num_heads
213 module.num_heads = query.size(2)
214 try:
215 output, softmax_max, softmax_sum = original_sparse(
216 module, query, key, value, None, **call_kwargs)
217 finally:
218 module.num_heads = old_heads
219 if module.training and not module.freeze_dsa:
220 _dsa_tensor_context.set(_DSATensorContext(
221 query=query, key=key, q_pe=q_pe, k_pe=k_pe))
222 output = _head_to_sequence(output, context)
223 if output.shape[:3] != torch.Size(local_shape[:3]):
224 raise RuntimeError(
225 f"DSA CP output {tuple(output.shape)} does not restore "
226 f"{local_shape}")
227 return output, softmax_max, softmax_sum
229 setattr(sparse_attention_with_gathered_kv, _ULYSSES_WRAPPED_FLAG, True)
230 attention_functions["dsa_sparse_attention"] = (
231 sparse_attention_with_gathered_kv)
233 if not getattr(original_kl, _ULYSSES_WRAPPED_FLAG, False):
234 class CPDSAKLLoss:
235 """Proxy the DSA KL loss with CP-transformed attention inputs."""
237 @staticmethod
238 def apply(*args: Any, **kwargs: Any) -> Any:
239 """Apply the original KL loss with saved global sequence tensors."""
240 arguments = _DSA_KL_SIGNATURE.bind(*args, **kwargs).arguments
241 saved = _dsa_tensor_context.get()
242 if saved is not None:
243 _dsa_tensor_context.set(None)
244 length = saved.query.size(1)
245 arguments.update(
246 query=saved.query.flatten(0, 1),
247 key=saved.key.flatten(0, 1),
248 query_rope=saved.q_pe.flatten(0, 1),
249 key_rope=saved.k_pe.flatten(0, 1),
250 actual_seq_qlen=_global_seq_len(
251 arguments["actual_seq_qlen"], length, saved.query.device),
252 actual_seq_klen=_global_seq_len(
253 arguments["actual_seq_klen"], length, saved.query.device),
254 )
255 return original_kl.apply(
256 *(arguments[name] for name in _DSA_KL_ARGUMENTS))
258 setattr(CPDSAKLLoss, _ULYSSES_WRAPPED_FLAG, True)
259 attention_module.SparseLightningIndexerKLLossTrainFunction = (
260 CPDSAKLLoss)
263def _head_tail_peer_rank(cp_mesh: Any) -> int:
264 """Return the mirror rank in the CP mesh."""
265 rank_list = list(cp_mesh.rank_list)
266 return rank_list[cp_mesh.size() - 1 - rank_list.index(dist.get_rank())]
269def _run_head_tail_half(
270 attention_fn: Callable[[Tensor, Tensor, Tensor, dict[str, Any]], Any],
271 query: Tensor, key: Tensor, value: Tensor,
272 call_kwargs: dict[str, Any]) -> Tensor:
273 """Run attention for one query half and check its output contract."""
274 output = attention_fn(query, key, value, call_kwargs)
275 if not isinstance(output, Tensor):
276 raise TypeError(
277 "Head-Tail load balance requires the attention callable to "
278 f"return a Tensor, got {type(output).__name__}"
279 )
280 return output
283def head_tail_load_balance_attention(
284 attention_fn: Callable[[Tensor, Tensor, Tensor, dict[str, Any]], Any],
285 query: Tensor, key: Tensor, value: Tensor,
286 attention_kwargs: dict[str, Any], cp_mesh: Any, *,
287 peer_attention_kwargs: Optional[dict[str, Any]] = None) -> Any:
288 """Run local-tensor Colossal Head-Tail communication.
290 The caller prepares split-specific mask or position metadata. When
291 ``peer_attention_kwargs`` is omitted, both attention calls reuse
292 ``attention_kwargs`` for backward compatibility.
293 """
294 if cp_mesh is None or cp_mesh.size() <= 1:
295 raise ValueError("Head-Tail load balance requires an active CP mesh")
296 local_q_len = query.shape[2]
297 if local_q_len % 2:
298 raise ValueError(
299 "Head-Tail load balance requires an even local Q sequence "
300 f"length, got {local_q_len}; pad the global sequence to a "
301 f"multiple of 2 * cp_size ({2 * cp_mesh.size()})"
302 )
304 peer_rank = _head_tail_peer_rank(cp_mesh)
305 query_peer = communication.p2p_exchange(
306 query.narrow(2, local_q_len // 2, local_q_len // 2), peer_rank)
307 global_key, global_value = flex_cp_allgather(key, value, 2, cp_mesh)
308 keep_output = _run_head_tail_half(
309 attention_fn, query.narrow(2, 0, local_q_len // 2),
310 global_key, global_value, attention_kwargs)
311 peer_output = _run_head_tail_half(
312 attention_fn, query_peer, global_key, global_value,
313 attention_kwargs if peer_attention_kwargs is None else peer_attention_kwargs,
314 )
315 return torch.cat(
316 [keep_output, communication.p2p_exchange(peer_output, peer_rank)], dim=2)
319def _cp_offset_causal_mask(q_len: int, kv_len: int, lo: int,
320 device, dtype=torch.bool):
321 """D-04: offset-aware causal mask (this rank's Q chunk has global offset lo).
323 Attendable positions: j <= lo + i (i is the local Q row index).
324 Replaces is_causal=True -- torch SDPA's is_causal is top-left aligned when
325 q_len != kv_len (equivalent to assuming Q starts at global position 0), so
326 under CP the chunks of rank>0 would be incorrectly masked (G4).
327 """
328 i = torch.arange(q_len, device=device).view(-1, 1)
329 j = torch.arange(kv_len, device=device).view(1, -1)
330 return (j <= (lo + i)).to(dtype)