Star137
193
代码
代码
Issues281
Pull Requests201
项目讨论
Wiki
分析
Star137
193
已合并
linear attention cp #1035
xu-xianliang创建于 7月13日
linear attention cp #1035
已合并
xu-xianliang创建于 7月13日
从xu-xianliang/hyper-parallellinear_cp合入到mindspore/hyper-parallelmaster
共 5 个文件变更+1568-68
Ahyper_parallel/core/context_parallel/linear_attention_context_parallel.py+1173-0
@@ -0,0 +1,1173 @@
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 execution for Qwen3.5-style Gated DeltaNet layers."""
16+from __future__ import annotations
17+ 
18+from typing import NamedTuple, Optional
19+ 
20+import torch
21+import torch.distributed as dist
22+from torch import nn
23+from torch.nn import functional as F
24+from torch.utils.checkpoint import checkpoint
25+ 
26+from hyper_parallel.core.context_parallel.context_parallel import (
27+ _ensure_1d,
28+)
29+from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
30+from hyper_parallel.core.dtensor.dtensor import DTensor
31+from hyper_parallel.core.tensor_parallel.style import ParallelStyle
32+from hyper_parallel.models.modules.linear_attention import torch_chunk_gated_delta_rule
33+from hyper_parallel.platform import get_platform
34+ 
35+ 
36+platform = get_platform()
37+ 
38+ 
39+def _global_peer_rank(cp_mesh: DeviceMesh, local_rank: int) -> int:
40+ """Map a CP-local rank index to its global distributed rank."""
41+ return int(cp_mesh.rank_list[local_rank])
42+ 
43+ 
44+def _slice_local_cp(
45+ tensor: torch.Tensor,
46+ dim: int,
47+ cp_rank: int,
48+ cp_size: int,
49+) -> torch.Tensor:
50+ """Return this CP rank's contiguous slice along ``dim``."""
51+ dim_size = tensor.shape[dim]
52+ if dim_size % cp_size != 0:
53+ raise ValueError(
54+ f"linear attention CP expects dim size {dim_size} "
55+ f"to be divisible by cp_size {cp_size}."
56+ )
57+ chunk = dim_size // cp_size
58+ return tensor.narrow(dim, cp_rank * chunk, chunk)
59+ 
60+ 
61+def _slice_qkv_local_cp(
62+ tensor: torch.Tensor,
63+ *,
64+ key_dim: int,
65+ value_dim: int,
66+ dim: int,
67+ cp_rank: int,
68+ cp_size: int,
69+) -> torch.Tensor:
70+ """Slice a fused ``[Q, K, V]`` tensor on the Q/K/V channel dimension."""
71+ q, k, v = torch.split(tensor, [key_dim, key_dim, value_dim], dim=dim)
72+ return torch.cat(
73+ (
74+ _slice_local_cp(q, dim, cp_rank, cp_size),
75+ _slice_local_cp(k, dim, cp_rank, cp_size),
76+ _slice_local_cp(v, dim, cp_rank, cp_size),
77+ ),
78+ dim=dim,
79+ )
80+ 
81+ 
82+def _local_tensor_at_cp_boundary(tensor: torch.Tensor) -> torch.Tensor:
83+ """Return the local tensor carried by a CP-boundary input.
84+ 
85+ The first supported Qwen3.5 linear-attention CP path keeps decoder-layer
86+ activations as local sequence shards. If an upstream wrapper passes that
87+ shard as a DTensor, use its local tensor and continue with the same
88+ ``[B, S_local, H]`` boundary contract.
89+ """
90+ if isinstance(tensor, DTensor):
91+ return tensor.to_local()
92+ return tensor
93+ 
94+ 
95+def _all_to_all_previous_rank_halo(
96+ tail: torch.Tensor,
97+ cp_mesh: DeviceMesh,
98+ cp_rank: int,
atomgit-bot
atomgit-botatomgit-bot7月13日

🔵 Low Priority

_all_gather_seq 函数签名包含 cp_rank: int 参数，但在函数体内完全没有使用该参数。该函数仅依赖 cp_mesh.rank_list 来确定 all-gather 的顺序。

调用处（_conv1d_with_all_gather_halo 第733行）仍传入 self.cp_rank，暗示调用者期望该参数有实际作用，但实际上它被静默忽略。这是一个死代码问题，可能在后续维护中造成误导。

changed line: _all_gather_seq 函数定义（第94-110行）中的 cp_rank 参数 → 参数声明但未使用 → 死代码 / 误导性 API。

建议：移除未使用的 cp_rank 参数，或将其用于实际逻辑（如根据 cp_rank 进行排序校验）。如果只是声明层面的干净清理，从签名和所有调用处删除即可。

改动建议
98
- cp_rank: int,
98
+ cp_size: int,
应用建议
likedislike
不准确？
99+ cp_size: int,
100+) -> torch.Tensor:
101+ """Send a convolution halo only to the next rank using differentiable A2AV."""
102+ if cp_size == 1:
103+ return torch.zeros_like(tail)
104+ 
105+ cp_group = cp_mesh.get_group()
106+ group_ranks = tuple(int(rank) for rank in dist.get_process_group_ranks(cp_group))
107+ rank_list = tuple(int(rank) for rank in cp_mesh.rank_list)
108+ rank_to_group_index = {rank: index for index, rank in enumerate(group_ranks)}
109+ halo_width = tail.shape[1]
110+ 
111+ input_splits = [0] * cp_size
112+ exchange_input = tail.permute(1, 0, 2).contiguous()
113+ if cp_rank < cp_size - 1:
114+ input_splits[rank_to_group_index[rank_list[cp_rank + 1]]] = halo_width
115+ else:
116+ exchange_input = exchange_input[:0]
117+ 
118+ output_splits = [0] * cp_size
119+ if cp_rank > 0:
120+ output_splits[rank_to_group_index[rank_list[cp_rank - 1]]] = halo_width
121+ 
122+ exchange_output = platform.differentiable_all_to_all_single(
123+ exchange_input,
124+ input_splits,
125+ output_splits,
126+ group=cp_group,
127+ )
128+ if cp_rank == 0:
129+ return torch.zeros_like(tail) + exchange_output.sum().to(tail.dtype) * 0
130+ return exchange_output.permute(1, 0, 2).contiguous()
131+ 
132+ 
133+def _causal_conv1d_with_cp_halo(
134+ mixed_qkv: torch.Tensor,
135+ conv1d: nn.Conv1d,
136+ cp_mesh: DeviceMesh,
137+ cp_rank: int,
138+ cp_size: int,
139+) -> torch.Tensor:
140+ """Run causal depthwise Conv1d with only the previous rank's boundary."""
141+ kernel_size = conv1d.kernel_size[0]
142+ dilation = conv1d.dilation[0]
143+ halo_width = (kernel_size - 1) * dilation
144+ if halo_width == 0 or cp_size == 1:
145+ conv_out = conv1d(mixed_qkv.transpose(1, 2))
146+ return F.silu(conv_out[:, :, : mixed_qkv.shape[1]]).transpose(1, 2)
147+ 
148+ if mixed_qkv.shape[1] < halo_width:
149+ raise ValueError(
150+ "linear attention CP conv halo requires local_seq_len >= "
151+ f"{halo_width}, got {mixed_qkv.shape[1]}."
152+ )
153+ 
154+ halo = _all_to_all_previous_rank_halo(
155+ mixed_qkv[:, -halo_width:, :].contiguous(),
156+ cp_mesh,
157+ cp_rank,
158+ cp_size,
159+ )
160+ conv_input = torch.cat((halo, mixed_qkv), dim=1).transpose(1, 2)
161+ conv_out = F.conv1d(
162+ input=conv_input,
163+ weight=conv1d.weight,
164+ bias=conv1d.bias,
165+ stride=conv1d.stride,
166+ padding=0,
167+ dilation=conv1d.dilation,
168+ groups=conv1d.groups,
169+ )
170+ return F.silu(conv_out).transpose(1, 2)
171+ 
172+ 
173+def _all_gather_stack(
174+ tensor: torch.Tensor,
175+ cp_mesh: DeviceMesh,
176+ cp_size: int,
177+) -> torch.Tensor:
178+ """All-gather equal-shaped tensors and stack them on a leading rank dim."""
179+ if cp_size == 1:
180+ return tensor.unsqueeze(0)
181+ return platform.differentiable_all_gather_concat(
182+ tensor.unsqueeze(0),
183+ cp_mesh.get_group(),
184+ cp_size,
185+ 0,
186+ tuple(int(rank) for rank in cp_mesh.rank_list),
187+ )
188+ 
189+ 
190+def _l2norm_torch(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor:
191+ """Match the pure torch GDN reference l2norm helper."""
192+ return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
193+ 
194+ 
195+class _GDNPreparedChunks(NamedTuple):
196+ """Reusable chunk intermediates shared by state-summary CP modes."""
197+ 
198+ initial_dtype: torch.dtype
199+ query: torch.Tensor
200+ key: torch.Tensor
201+ chunk_value: torch.Tensor
202+ g: torch.Tensor
203+ decay_mask: torch.Tensor
204+ k_cumdecay: torch.Tensor
205+ sequence_length: int
206+ total_sequence_length: int
207+ chunk_size: int
208+ 
209+ 
210+def _prepare_gdn_chunks_for_summary(
211+ query: torch.Tensor,
212+ key: torch.Tensor,
213+ value: torch.Tensor,
214+ g: torch.Tensor,
215+ beta: torch.Tensor,
216+ *,
217+ chunk_size: int = 64,
218+ use_qk_l2norm_in_kernel: bool = False,
219+) -> _GDNPreparedChunks:
220+ """Prepare GDN chunk intermediates shared by summary and local output."""
221+ initial_dtype = query.dtype
222+ if use_qk_l2norm_in_kernel:
223+ query = _l2norm_torch(query, dim=-1, eps=1e-6)
224+ key = _l2norm_torch(key, dim=-1, eps=1e-6)
225+ 
226+ query, key, value, beta, g = [
227+ x.transpose(1, 2).contiguous().to(torch.float32)
228+ for x in (query, key, value, beta, g)
229+ ]
230+ 
231+ sequence_length = key.shape[2]
232+ pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
233+ query = F.pad(query, (0, 0, 0, pad_size))
234+ key = F.pad(key, (0, 0, 0, pad_size))
235+ value = F.pad(value, (0, 0, 0, pad_size))
236+ beta = F.pad(beta, (0, pad_size))
237+ g = F.pad(g, (0, pad_size))
238+ total_sequence_length = sequence_length + pad_size
239+ 
240+ query = query * (1 / (query.shape[-1] ** 0.5))
241+ v_beta = value * beta.unsqueeze(-1)
242+ k_beta = key * beta.unsqueeze(-1)
243+ query, key, k_beta, v_beta = [
244+ x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
245+ for x in (query, key, k_beta, v_beta)
246+ ]
247+ g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
248+ 
249+ mask = torch.triu(
250+ torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
251+ diagonal=0,
252+ )
253+ g = g.cumsum(dim=-1)
254+ decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
255+ attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
256+ for row_idx in range(1, chunk_size):
257+ row = attn[..., row_idx, :row_idx].clone()
258+ sub = attn[..., :row_idx, :row_idx].clone()
259+ attn[..., row_idx, :row_idx] = row + (row.unsqueeze(-1) * sub).sum(-2)
260+ attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
261+ 
262+ chunk_value = attn @ v_beta
263+ k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
264+ return _GDNPreparedChunks(
265+ initial_dtype=initial_dtype,
266+ query=query,
267+ key=key,
268+ chunk_value=chunk_value,
269+ g=g,
270+ decay_mask=decay_mask,
271+ k_cumdecay=k_cumdecay,
272+ sequence_length=sequence_length,
273+ total_sequence_length=total_sequence_length,
274+ chunk_size=chunk_size,
275+ )
276+ 
277+ 
278+def _compute_gdn_state_summary_from_prepared(
279+ prepared: _GDNPreparedChunks,
280+) -> tuple[torch.Tensor, torch.Tensor]:
281+ """Compute ``state_out = M @ state_in + S`` from prepared GDN chunks."""
282+ key = prepared.key
283+ batch_size, num_heads, _, _, k_head_dim = key.shape
284+ v_head_dim = prepared.chunk_value.shape[-1]
285+ eye = torch.eye(k_head_dim, device=key.device, dtype=torch.float32).reshape(
286+ 1, 1, k_head_dim, k_head_dim
287+ )
288+ state_ext = torch.zeros(
289+ batch_size,
290+ num_heads,
291+ k_head_dim,
292+ v_head_dim,
293+ device=key.device,
294+ dtype=torch.float32,
295+ )
296+ transition = eye.expand(batch_size, num_heads, -1, -1).clone()
297+ 
298+ for chunk_idx in range(key.shape[2]):
299+ key_i = key[:, :, chunk_idx]
300+ value_i = prepared.chunk_value[:, :, chunk_idx]
301+ w_i = prepared.k_cumdecay[:, :, chunk_idx]
302+ g_i = prepared.g[:, :, chunk_idx]
303+ decay = g_i[:, :, -1].exp()
304+ key_decay = key_i * (g_i[:, :, -1, None] - g_i).exp()[..., None]
305+ 
306+ transition_i = (
307+ decay[:, :, None, None] * eye
308+ - key_decay.transpose(-1, -2) @ w_i
309+ )
310+ state_ext_i = key_decay.transpose(-1, -2) @ value_i
311+ state_ext = transition_i @ state_ext + state_ext_i
312+ transition = transition_i @ transition
313+ 
314+ return state_ext, transition
315+ 
316+ 
317+def _checkpoint_gdn_state_summary(
318+ prepared: _GDNPreparedChunks,
319+) -> tuple[torch.Tensor, torch.Tensor]:
320+ """Compute a state summary without retaining its per-chunk autograd graph."""
321+ if not torch.is_grad_enabled():
322+ return _compute_gdn_state_summary_from_prepared(prepared)
323+ 
324+ def recompute(
325+ key: torch.Tensor,
326+ chunk_value: torch.Tensor,
327+ g: torch.Tensor,
328+ k_cumdecay: torch.Tensor,
329+ ) -> tuple[torch.Tensor, torch.Tensor]:
330+ """Rebuild a prepared view from explicit checkpoint inputs."""
331+ checkpoint_prepared = prepared._replace(
332+ key=key,
333+ chunk_value=chunk_value,
334+ g=g,
335+ k_cumdecay=k_cumdecay,
336+ )
337+ return _compute_gdn_state_summary_from_prepared(checkpoint_prepared)
338+ 
339+ return checkpoint(
340+ recompute,
341+ prepared.key,
342+ prepared.chunk_value,
343+ prepared.g,
344+ prepared.k_cumdecay,
345+ use_reentrant=False,
346+ preserve_rng_state=False,
347+ )
348+ 
349+ 
350+def _run_prepared_gdn_chunks(
351+ prepared: _GDNPreparedChunks,
352+ initial_state: Optional[torch.Tensor],
353+) -> torch.Tensor:
354+ """Run local GDN output using already prepared chunk intermediates."""
355+ query = prepared.query
356+ key = prepared.key
357+ chunk_value = prepared.chunk_value
358+ batch_size, num_heads, _, _, k_head_dim = key.shape
359+ v_head_dim = chunk_value.shape[-1]
360+ recurrent_state = (
361+ torch.zeros(
362+ batch_size,
363+ num_heads,
364+ k_head_dim,
365+ v_head_dim,
366+ device=chunk_value.device,
367+ dtype=chunk_value.dtype,
368+ )
369+ if initial_state is None
370+ else initial_state.to(chunk_value)
371+ )
372+ core_attn_out = torch.zeros_like(chunk_value)
373+ 
374+ for chunk_idx in range(0, prepared.total_sequence_length // prepared.chunk_size):
375+ q_i = query[:, :, chunk_idx]
376+ k_i = key[:, :, chunk_idx]
377+ v_i = chunk_value[:, :, chunk_idx]
378+ attn = q_i @ k_i.transpose(-1, -2) * prepared.decay_mask[:, :, chunk_idx]
379+ v_prime = prepared.k_cumdecay[:, :, chunk_idx] @ recurrent_state
380+ v_new = v_i - v_prime
381+ attn_inter = (
382+ q_i * prepared.g[:, :, chunk_idx, :, None].exp()
383+ ) @ recurrent_state
384+ core_attn_out[:, :, chunk_idx] = attn_inter + attn @ v_new
385+ recurrent_state = (
386+ recurrent_state * prepared.g[:, :, chunk_idx, -1, None, None].exp()
387+ + (
388+ k_i
389+ * (
390+ prepared.g[:, :, chunk_idx, -1, None]
391+ - prepared.g[:, :, chunk_idx]
392+ ).exp()[..., None]
393+ ).transpose(-1, -2) @ v_new
394+ )
395+ 
396+ core_attn_out = core_attn_out.reshape(
397+ core_attn_out.shape[0],
398+ core_attn_out.shape[1],
399+ -1,
400+ core_attn_out.shape[-1],
401+ )
402+ core_attn_out = core_attn_out[:, :, :prepared.sequence_length]
403+ return core_attn_out.transpose(1, 2).contiguous().to(prepared.initial_dtype)
404+ 
405+ 
406+def _pack_gdn_state_summary(
407+ state_ext: torch.Tensor,
408+ transition: torch.Tensor,
409+) -> torch.Tensor:
410+ """Pack ``S`` and ``M`` summaries into one all-gather payload."""
411+ if state_ext.shape[:-1] != transition.shape[:-1]:
412+ raise ValueError(
413+ "state_ext and transition must share [B,H,K] dimensions, got "
414+ f"{tuple(state_ext.shape)} and {tuple(transition.shape)}."
415+ )
416+ return torch.cat((state_ext, transition), dim=-1)
417+ 
418+ 
419+def _unpack_gdn_state_summary(
420+ packed: torch.Tensor,
421+ v_head_dim: int,
422+) -> tuple[torch.Tensor, torch.Tensor]:
423+ """Unpack a gathered ``[S, M]`` payload."""
424+ if packed.shape[-1] <= v_head_dim:
425+ raise ValueError(
426+ f"packed state summary last dim must be > v_head_dim={v_head_dim}, "
427+ f"got {packed.shape[-1]}."
428+ )
429+ state_ext = packed[..., :v_head_dim]
430+ transition = packed[..., v_head_dim:]
431+ return state_ext, transition
432+ 
433+ 
434+def _merge_gdn_prefix_state_summaries_torch(
435+ state_ext: torch.Tensor,
436+ transition: torch.Tensor,
437+ rank: int,
438+) -> torch.Tensor:
439+ """Merge gathered GDN summaries before ``rank`` into its initial state."""
440+ if state_ext.dim() != 5 or transition.dim() != 5:
441+ raise ValueError(
442+ "state summary merge expects state_ext [R,B,H,K,V] and "
443+ "transition [R,B,H,K,K]."
444+ )
445+ if state_ext.shape[0] != transition.shape[0]:
446+ raise ValueError("state_ext and transition must have the same rank dimension.")
447+ if rank < 0 or rank > state_ext.shape[0]:
448+ raise ValueError(f"rank must be in [0, {state_ext.shape[0]}], got {rank}.")
449+ 
450+ state = torch.zeros_like(state_ext[0])
451+ for prev_rank in range(rank):
452+ state = transition[prev_rank] @ state + state_ext[prev_rank]
453+ return state
454+ 
455+ 
456+def _gdn_state_all_gather(
457+ query: torch.Tensor,
458+ key: torch.Tensor,
459+ value: torch.Tensor,
460+ g: torch.Tensor,
461+ beta: torch.Tensor,
462+ cp_mesh: DeviceMesh,
463+ cp_rank: int,
464+ cp_size: int,
465+ *,
466+ use_qk_l2norm_in_kernel: bool,
467+) -> torch.Tensor:
468+ """Apply local GDN with all-gathered recurrent-state summaries."""
469+ if cp_size == 1:
470+ core_attn_out, _ = torch_chunk_gated_delta_rule(
471+ query,
472+ key,
473+ value,
474+ g=g,
475+ beta=beta,
476+ initial_state=None,
477+ output_final_state=False,
478+ use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
479+ )
480+ return core_attn_out
481+ 
482+ prepared = _prepare_gdn_chunks_for_summary(
483+ query,
484+ key,
485+ value,
486+ g,
487+ beta,
488+ use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
489+ )
490+ state_ext, transition = _checkpoint_gdn_state_summary(prepared)
491+ packed_summary = _pack_gdn_state_summary(state_ext, transition)
492+ gathered_summary = _all_gather_stack(packed_summary, cp_mesh, cp_size)
493+ gathered_state_ext, gathered_transition = _unpack_gdn_state_summary(
494+ gathered_summary,
495+ state_ext.shape[-1],
496+ )
497+ 
498+ initial_state = _merge_gdn_prefix_state_summaries_torch(
499+ gathered_state_ext,
500+ gathered_transition,
501+ cp_rank,
502+ )
503+ all_gather_tie = gathered_summary.sum()
504+ initial_state = initial_state + all_gather_tie.to(initial_state.dtype) * 0
505+ return _run_prepared_gdn_chunks(prepared, initial_state)
506+ 
507+ 
508+class _RecvInitialStateP2PFunction(torch.autograd.Function):
509+ """Receive the recurrent initial state; send its gradient in backward."""
510+ 
511+ @staticmethod
512+ def forward( # pylint: disable=arguments-differ
513+ ctx,
司小南（机器人）
司小南（机器人）司小南（机器人）7月13日

此条代码评论区间+511至+513

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。

likedislike
514+ anchor: torch.Tensor,
515+ cp_group,
516+ prev_rank: int,
517+ state_shape: tuple[int, ...],
518+ ) -> torch.Tensor:
519+ """Receive the initial state from the preceding CP rank."""
520+ state = torch.empty(state_shape, device=anchor.device, dtype=torch.float32)
521+ dist.recv(state, src=prev_rank, group=cp_group)
522+ ctx.cp_group = cp_group
523+ ctx.prev_rank = prev_rank
524+ return state
525+ 
526+ @staticmethod
527+ def backward(ctx, grad_state: Optional[torch.Tensor]):
528+ if grad_state is None:
529+ raise RuntimeError("linear attention P2P backward missing initial-state grad.")
530+ dist.send(grad_state.contiguous(), dst=ctx.prev_rank, group=ctx.cp_group)
531+ return None, None, None, None
532+ 
533+ 
534+class _SendFinalStateP2PFunction(torch.autograd.Function):
535+ """Send the recurrent final state; receive its gradient in backward."""
536+ 
537+ @staticmethod
538+ def forward( # pylint: disable=arguments-differ
539+ ctx,
司小南（机器人）
司小南（机器人）司小南（机器人）7月13日

此条代码评论区间+537至+539

【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：pylint，请Committer检视其合理性。

likedislike
540+ final_state: torch.Tensor,
541+ cp_group,
542+ next_rank: int,
543+ ) -> torch.Tensor:
544+ """Send the final state to the succeeding CP rank."""
545+ dist.send(final_state.contiguous(), dst=next_rank, group=cp_group)
546+ ctx.cp_group = cp_group
547+ ctx.next_rank = next_rank
548+ ctx.state_shape = tuple(final_state.shape)
549+ ctx.state_dtype = final_state.dtype
550+ return final_state.new_zeros(())
551+ 
552+ @staticmethod
553+ def backward(ctx, grad_token: torch.Tensor):
554+ grad_state = torch.empty(
555+ ctx.state_shape,
556+ device=grad_token.device,
557+ dtype=ctx.state_dtype,
558+ )
559+ dist.recv(grad_state, src=ctx.next_rank, group=ctx.cp_group)
560+ return grad_state, None, None
561+ 
562+ 
563+def _apply_gdn_state_summary(
564+ state_ext: torch.Tensor,
565+ transition: torch.Tensor,
566+ initial_state: Optional[torch.Tensor],
567+) -> torch.Tensor:
568+ """Apply ``state_out = M @ state_in + S`` to an incoming GDN state."""
569+ if initial_state is None:
570+ return state_ext
571+ return transition @ initial_state.to(transition) + state_ext
572+ 
573+ 
574+def _gdn_state_p2p_summary(
575+ query: torch.Tensor,
576+ key: torch.Tensor,
577+ value: torch.Tensor,
578+ g: torch.Tensor,
579+ beta: torch.Tensor,
580+ cp_mesh: DeviceMesh,
581+ cp_rank: int,
582+ cp_size: int,
583+ *,
584+ use_qk_l2norm_in_kernel: bool,
585+) -> torch.Tensor:
586+ """Run local GDN with an affine-summary state wavefront.
587+ 
588+ Every rank prepares its local chunks and state transition in parallel.
589+ The rank-ordered critical path then contains only ``M @ state + S`` and
590+ the small state transfer. Token outputs retain the ordinary PyTorch graph,
591+ while the two custom autograd boundaries reverse the state communication.
592+ """
593+ if cp_size == 1:
594+ core_attn_out, _ = torch_chunk_gated_delta_rule(
595+ query,
596+ key,
597+ value,
598+ g=g,
599+ beta=beta,
600+ initial_state=None,
601+ output_final_state=False,
602+ use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
603+ )
604+ return core_attn_out
605+ 
606+ cp_group = cp_mesh.get_group()
607+ prev_rank = _global_peer_rank(cp_mesh, cp_rank - 1) if cp_rank > 0 else -1
608+ next_rank = _global_peer_rank(cp_mesh, cp_rank + 1) if cp_rank < cp_size - 1 else -1
609+ state_shape = (query.shape[0], value.shape[2], query.shape[3], value.shape[3])
610+ 
611+ prepared = _prepare_gdn_chunks_for_summary(
612+ query,
613+ key,
614+ value,
615+ g,
616+ beta,
617+ use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
618+ )
619+ state_ext = None
620+ transition = None
621+ if cp_rank < cp_size - 1:
622+ state_ext, transition = _checkpoint_gdn_state_summary(prepared)
623+ 
624+ initial_state = None
625+ if cp_rank > 0:
626+ initial_state = _RecvInitialStateP2PFunction.apply(
627+ query,
628+ cp_group,
629+ prev_rank,
630+ state_shape,
631+ )
632+ 
633+ send_token = None
634+ if cp_rank < cp_size - 1:
635+ final_state = _apply_gdn_state_summary(
636+ state_ext,
637+ transition,
638+ initial_state,
639+ )
640+ send_token = _SendFinalStateP2PFunction.apply(final_state, cp_group, next_rank)
641+ 
642+ core_attn_out = _run_prepared_gdn_chunks(prepared, initial_state)
643+ if send_token is not None:
644+ core_attn_out = core_attn_out + send_token.to(core_attn_out.dtype) * 0
645+ 
646+ return core_attn_out
647+ 
648+ 
649+def _differentiable_all_to_all_shard(
650+ tensor: torch.Tensor,
651+ device_mesh: DeviceMesh,
652+ *,
653+ split_dim: int,
654+ concat_dim: int,
655+) -> torch.Tensor:
656+ """Split local data on ``split_dim`` and concatenate peers on ``concat_dim``.
657+ 
658+ This is the local-tensor equivalent of DTensor ``Shard(concat_dim) ->
659+ Shard(split_dim)`` redistribution for a 1-D mesh. It uses platform-level
660+ differentiable all-to-all directly to avoid wrapping each activation in a
661+ temporary DTensor.
662+ """
663+ split_count = device_mesh.size()
664+ if split_count == 1:
665+ return tensor
666+ 
667+ original_shape = tuple(tensor.shape)
668+ dim_size = original_shape[split_dim]
669+ if dim_size % split_count != 0:
670+ raise ValueError(
671+ f"linear attention all-to-all split dim {split_dim} with size "
672+ f"{dim_size} must be divisible by cp_size {split_count}."
673+ )
674+ 
675+ split_size = dim_size // split_count
676+ final_shape = list(original_shape)
677+ if split_dim != concat_dim:
678+ final_shape[split_dim] = split_size
679+ final_shape[concat_dim] = final_shape[concat_dim] * split_count
680+ final_shape = tuple(final_shape)
681+ 
682+ reshape_dims = list(original_shape)
683+ reshape_dims[split_dim] = split_count
684+ reshape_dims.insert(split_dim + 1, split_size)
685+ 
686+ trans_dims = list(range(len(reshape_dims)))
687+ trans_dims.remove(split_dim)
688+ trans_dims.insert(0, split_dim)
689+ 
690+ a2a_input = tensor.reshape(reshape_dims).permute(trans_dims).contiguous()
691+ reshape_shape = list(a2a_input.shape)
692+ reshape_shape[0] = reshape_shape[0] * reshape_shape[1]
693+ reshape_shape.pop(1)
694+ a2a_input = a2a_input.reshape(reshape_shape)
695+ 
696+ a2a_input = a2a_input.contiguous()
697+ split_len = a2a_input.shape[0] // split_count
698+ input_splits = [split_len] * split_count
699+ output_splits = [split_len] * split_count
700+ output = platform.differentiable_all_to_all_single(
701+ a2a_input,
702+ input_splits,
703+ output_splits,
704+ group=device_mesh.get_group(),
705+ )
706+ 
707+ output_reshape = list(output.shape)
708+ output_reshape[0] = split_count
709+ output_reshape.insert(1, output.shape[0] // split_count)
710+ 
711+ out_trans_dims = list(range(len(output_reshape)))
712+ first_dim = out_trans_dims.pop(0)
713+ if concat_dim >= len(out_trans_dims):
714+ out_trans_dims.append(first_dim)
715+ else:
716+ out_trans_dims.insert(concat_dim, first_dim)
717+ 
718+ final_output = output.reshape(output_reshape).permute(out_trans_dims).contiguous()
719+ final_reshape = list(final_output.shape)
720+ if concat_dim < len(final_reshape) - 1:
721+ final_reshape[concat_dim] = (
722+ final_reshape[concat_dim] * final_reshape[concat_dim + 1]
723+ )
724+ final_reshape.pop(concat_dim + 1)
725+ 
726+ return final_output.reshape(final_reshape).view(final_shape)
727+ 
728+ 
729+class LinearAttentionUlyssesCPWrapper(nn.Module):
730+ """Pure-Ulysses CP execution wrapper for a Qwen3.5 Gated DeltaNet module.
731+ 
732+ Parameters stay owned by the original module. The wrapper only changes the
733+ execution layout:
734+ 
735+ ``[B, S_local, full_heads] -> [B, S_full, local_heads] ->
736+ [B, S_local, full_heads]``.
737+ """
738+ 
739+ def __init__(self, module: nn.Module, device_mesh: DeviceMesh):
740+ super().__init__()
741+ self.module = module
742+ self.cp_mesh = _ensure_1d(device_mesh)
743+ self.cp_size = self.cp_mesh.size()
744+ self.cp_rank = self.cp_mesh.get_local_rank()
745+ self.seq_dim = 1
746+ self.head_dim = 2
747+ self._validate_module()
748+ 
749+ def _validate_module(self) -> None:
750+ if self.cp_size <= 1:
751+ return
752+ if self.module.num_k_heads % self.cp_size != 0:
753+ raise ValueError(
754+ f"linear attention num_k_heads ({self.module.num_k_heads}) must be "
755+ f"divisible by cp_size ({self.cp_size}) for Ulysses CP."
756+ )
757+ if self.module.num_v_heads % self.cp_size != 0:
758+ raise ValueError(
759+ f"linear attention num_v_heads ({self.module.num_v_heads}) must be "
760+ f"divisible by cp_size ({self.cp_size}) for Ulysses CP."
761+ )
762+ 
763+ def _seq_to_head(self, tensor: torch.Tensor) -> torch.Tensor:
764+ return _differentiable_all_to_all_shard(
765+ tensor,
766+ self.cp_mesh,
767+ split_dim=self.head_dim,
768+ concat_dim=self.seq_dim,
769+ )
770+ 
771+ def _head_to_seq(self, tensor: torch.Tensor) -> torch.Tensor:
772+ return _differentiable_all_to_all_shard(
773+ tensor,
774+ self.cp_mesh,
775+ split_dim=self.seq_dim,
776+ concat_dim=self.head_dim,
777+ )
778+ 
779+ def _seq_to_head_qkvba(
780+ self,
781+ q_proj: torch.Tensor,
782+ k_proj: torch.Tensor,
783+ v_proj: torch.Tensor,
784+ b: torch.Tensor,
785+ a: torch.Tensor,
786+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
787+ """Pack Q/K/V/B/A by CP rank and run a single seq-to-head all-to-all."""
788+ if self.cp_size == 1:
789+ return q_proj, k_proj, v_proj, b, a
790+ 
791+ base = self.module
792+ local_key_dim = base.key_dim // self.cp_size
793+ local_value_dim = base.value_dim // self.cp_size
794+ local_num_v_heads = base.num_v_heads // self.cp_size
795+ 
796+ q_chunks = torch.split(q_proj, local_key_dim, dim=-1)
797+ k_chunks = torch.split(k_proj, local_key_dim, dim=-1)
798+ v_chunks = torch.split(v_proj, local_value_dim, dim=-1)
799+ b_chunks = torch.split(b, local_num_v_heads, dim=-1)
800+ a_chunks = torch.split(a, local_num_v_heads, dim=-1)
801+ rank_major_chunks = [
802+ torch.cat(chunks, dim=-1)
803+ for chunks in zip(q_chunks, k_chunks, v_chunks, b_chunks, a_chunks)
804+ ]
805+ packed = torch.cat(rank_major_chunks, dim=-1).contiguous()
806+ packed = self._seq_to_head(packed)
807+ return torch.split(
808+ packed,
809+ [
810+ local_key_dim,
811+ local_key_dim,
812+ local_value_dim,
813+ local_num_v_heads,
814+ local_num_v_heads,
815+ ],
816+ dim=-1,
817+ )
818+ 
819+ def _local_conv_weight(self) -> torch.Tensor:
820+ return _slice_qkv_local_cp(
821+ self.module.conv1d.weight,
822+ key_dim=self.module.key_dim,
823+ value_dim=self.module.value_dim,
824+ dim=0,
825+ cp_rank=self.cp_rank,
826+ cp_size=self.cp_size,
827+ )
828+ 
829+ def _local_conv_bias(self) -> Optional[torch.Tensor]:
830+ bias = self.module.conv1d.bias
831+ if bias is None:
832+ return None
833+ return _slice_qkv_local_cp(
834+ bias,
835+ key_dim=self.module.key_dim,
836+ value_dim=self.module.value_dim,
837+ dim=0,
838+ cp_rank=self.cp_rank,
839+ cp_size=self.cp_size,
840+ )
841+ 
842+ def forward(
843+ self,
844+ hidden_states: torch.Tensor,
845+ attention_mask: Optional[torch.Tensor] = None,
846+ **kwargs,
847+ ) -> torch.Tensor:
848+ """Run Gated DeltaNet with pure Ulysses context parallel."""
849+ del kwargs
850+ hidden_states = _local_tensor_at_cp_boundary(hidden_states)
851+ 
852+ base = self.module
853+ if attention_mask is not None and attention_mask.ndim == 2:
854+ hidden_states = hidden_states * attention_mask[:, :, None].to(
855+ hidden_states.dtype
856+ )
857+ 
858+ bsz, local_seq_len, _ = hidden_states.shape
859+ mixed_qkv = base.in_proj_qkv(hidden_states)
860+ z = base.in_proj_z(hidden_states).reshape(
861+ bsz,
862+ local_seq_len,
863+ base.num_v_heads,
864+ base.head_v_dim,
865+ )
866+ b = base.in_proj_b(hidden_states)
867+ a = base.in_proj_a(hidden_states)
868+ 
869+ q_proj, k_proj, v_proj = torch.split(
870+ mixed_qkv,
871+ [base.key_dim, base.key_dim, base.value_dim],
872+ dim=-1,
873+ )
874+ q_proj, k_proj, v_proj, b, a = self._seq_to_head_qkvba(q_proj, k_proj, v_proj, b, a)
875+ 
876+ full_seq_len = q_proj.shape[1]
877+ local_key_dim = base.key_dim // self.cp_size
878+ local_value_dim = base.value_dim // self.cp_size
879+ local_num_k_heads = base.num_k_heads // self.cp_size
880+ local_num_v_heads = base.num_v_heads // self.cp_size
881+ local_conv_dim = local_key_dim * 2 + local_value_dim
882+ 
883+ mixed_qkv = torch.cat((q_proj, k_proj, v_proj), dim=-1).transpose(1, 2)
884+ conv_out = F.conv1d(
885+ input=mixed_qkv,
886+ weight=self._local_conv_weight(),
887+ bias=self._local_conv_bias(),
888+ stride=base.conv1d.stride,
889+ padding=base.conv1d.padding,
890+ dilation=base.conv1d.dilation,
891+ groups=local_conv_dim,
892+ )
893+ mixed_qkv = F.silu(conv_out[:, :, :full_seq_len]).transpose(1, 2)
894+ 
895+ query, key, value = torch.split(
896+ mixed_qkv,
897+ [local_key_dim, local_key_dim, local_value_dim],
898+ dim=-1,
899+ )
900+ query = query.reshape(bsz, full_seq_len, local_num_k_heads, base.head_k_dim)
901+ key = key.reshape(bsz, full_seq_len, local_num_k_heads, base.head_k_dim)
902+ value = value.reshape(bsz, full_seq_len, local_num_v_heads, base.head_v_dim)
903+ 
904+ a_log = _slice_local_cp(base.A_log, 0, self.cp_rank, self.cp_size)
905+ dt_bias = _slice_local_cp(base.dt_bias, 0, self.cp_rank, self.cp_size)
906+ beta = b.sigmoid()
907+ g = -a_log.float().exp() * F.softplus(a.float() + dt_bias)
908+ 
909+ if base.kv_groups > 1:
910+ query = query.repeat_interleave(base.kv_groups, dim=2)
911+ key = key.repeat_interleave(base.kv_groups, dim=2)
912+ 
913+ core_attn_out, _ = torch_chunk_gated_delta_rule(
914+ query,
915+ key,
916+ value,
917+ g=g,
918+ beta=beta,
919+ initial_state=None,
920+ output_final_state=False,
921+ use_qk_l2norm_in_kernel=True,
922+ )
923+ 
924+ core_attn_out = self._head_to_seq(core_attn_out)
925+ core_attn_out = core_attn_out.reshape(-1, base.head_v_dim)
926+ z_flat = z.reshape(-1, base.head_v_dim)
927+ core_attn_out = base.norm(core_attn_out, z_flat)
928+ core_attn_out = core_attn_out.reshape(bsz, local_seq_len, base.value_dim)
929+ if hasattr(base, "out_proj_input"):
930+ core_attn_out = base.out_proj_input(core_attn_out)
931+ return base.out_proj(core_attn_out)
932+ 
933+ 
934+class LinearAttentionP2PCPWrapper(nn.Module):
935+ """Sequence-sharded GDN CP with an affine-summary state wavefront."""
936+ 
937+ def __init__(self, module: nn.Module, device_mesh: DeviceMesh):
938+ super().__init__()
939+ self.module = module
940+ self.cp_mesh = _ensure_1d(device_mesh)
941+ self.cp_size = self.cp_mesh.size()
942+ self.cp_rank = self.cp_mesh.get_local_rank()
943+ self._validate_module()
944+ 
945+ def _validate_module(self) -> None:
946+ """Validate the Conv1d requirements of the P2P CP path."""
947+ conv = self.module.conv1d
948+ if conv.stride != (1,):
949+ raise ValueError(
950+ "linear attention P2P CP currently supports only conv1d stride=1."
951+ )
952+ if conv.groups != self.module.conv_dim:
953+ raise ValueError(
954+ "linear attention P2P CP expects depthwise conv1d groups=conv_dim."
955+ )
956+ if (
957+ conv.in_channels != self.module.conv_dim
958+ or conv.out_channels != self.module.conv_dim
959+ ):
960+ raise ValueError(
961+ "linear attention P2P CP expects conv1d channels to match conv_dim."
962+ )
963+ 
964+ def _conv1d_with_halo(self, mixed_qkv: torch.Tensor) -> torch.Tensor:
965+ """Run local Conv1d after exchanging only the previous-rank halo."""
966+ return _causal_conv1d_with_cp_halo(
967+ mixed_qkv,
968+ self.module.conv1d,
969+ self.cp_mesh,
970+ self.cp_rank,
971+ self.cp_size,
972+ )
973+ 
974+ def forward(
975+ self,
976+ hidden_states: torch.Tensor,
977+ attention_mask: Optional[torch.Tensor] = None,
978+ **kwargs,
979+ ) -> torch.Tensor:
980+ """Run Gated DeltaNet on local sequence shards with recurrent-state P2P."""
981+ del kwargs
982+ hidden_states = _local_tensor_at_cp_boundary(hidden_states)
983+ 
984+ base = self.module
985+ if attention_mask is not None and attention_mask.ndim == 2:
986+ hidden_states = hidden_states * attention_mask[:, :, None].to(
987+ hidden_states.dtype
988+ )
989+ 
990+ bsz, local_seq_len, _ = hidden_states.shape
991+ mixed_qkv = base.in_proj_qkv(hidden_states)
992+ z = base.in_proj_z(hidden_states).reshape(
993+ bsz,
994+ local_seq_len,
995+ base.num_v_heads,
996+ base.head_v_dim,
997+ )
998+ b = base.in_proj_b(hidden_states)
999+ a = base.in_proj_a(hidden_states)
1000+ 
1001+ mixed_qkv = self._conv1d_with_halo(mixed_qkv)
1002+ query, key, value = torch.split(
1003+ mixed_qkv,
1004+ [base.key_dim, base.key_dim, base.value_dim],
1005+ dim=-1,
1006+ )
1007+ query = query.reshape(bsz, local_seq_len, base.num_k_heads, base.head_k_dim)
1008+ key = key.reshape(bsz, local_seq_len, base.num_k_heads, base.head_k_dim)
1009+ value = value.reshape(bsz, local_seq_len, base.num_v_heads, base.head_v_dim)
1010+ 
1011+ beta = b.sigmoid()
1012+ g = -base.A_log.float().exp() * F.softplus(a.float() + base.dt_bias)
1013+ 
1014+ if base.kv_groups > 1:
1015+ query = query.repeat_interleave(base.kv_groups, dim=2)
1016+ key = key.repeat_interleave(base.kv_groups, dim=2)
1017+ 
1018+ core_attn_out = _gdn_state_p2p_summary(
1019+ query,
1020+ key,
1021+ value,
1022+ g,
1023+ beta,
1024+ self.cp_mesh,
1025+ self.cp_rank,
1026+ self.cp_size,
1027+ use_qk_l2norm_in_kernel=True,
1028+ )
1029+ 
1030+ core_attn_out = core_attn_out.reshape(-1, base.head_v_dim)
1031+ z_flat = z.reshape(-1, base.head_v_dim)
1032+ core_attn_out = base.norm(core_attn_out, z_flat)
1033+ core_attn_out = core_attn_out.reshape(bsz, local_seq_len, base.value_dim)
1034+ if hasattr(base, "out_proj_input"):
1035+ core_attn_out = base.out_proj_input(core_attn_out)
1036+ return base.out_proj(core_attn_out)
1037+ 
1038+ 
1039+class LinearAttentionAllGatherCPWrapper(nn.Module):
1040+ """Sequence-sharded GDN CP using all-gathered recurrent-state summaries."""
1041+ 
1042+ def __init__(self, module: nn.Module, device_mesh: DeviceMesh):
1043+ super().__init__()
1044+ self.module = module
1045+ self.cp_mesh = _ensure_1d(device_mesh)
1046+ self.cp_size = self.cp_mesh.size()
1047+ self.cp_rank = self.cp_mesh.get_local_rank()
1048+ self._validate_module()
1049+ 
1050+ def _validate_module(self) -> None:
1051+ """Validate the Conv1d requirements of the all-gather CP path."""
1052+ conv = self.module.conv1d
1053+ if conv.stride != (1,):
1054+ raise ValueError(
1055+ "linear attention all-gather CP currently supports only "
1056+ "conv1d stride=1."
1057+ )
1058+ if conv.groups != self.module.conv_dim:
1059+ raise ValueError(
1060+ "linear attention all-gather CP expects depthwise conv1d "
1061+ "groups=conv_dim."
1062+ )
1063+ if (
1064+ conv.in_channels != self.module.conv_dim
1065+ or conv.out_channels != self.module.conv_dim
1066+ ):
1067+ raise ValueError(
1068+ "linear attention all-gather CP expects conv1d channels to "
1069+ "match conv_dim."
1070+ )
1071+ 
1072+ def _conv1d_with_halo(self, mixed_qkv: torch.Tensor) -> torch.Tensor:
1073+ """Run local Conv1d after exchanging only the previous-rank halo."""
1074+ return _causal_conv1d_with_cp_halo(
1075+ mixed_qkv,
1076+ self.module.conv1d,
1077+ self.cp_mesh,
1078+ self.cp_rank,
1079+ self.cp_size,
1080+ )
1081+ 
1082+ def forward(
1083+ self,
1084+ hidden_states: torch.Tensor,
1085+ attention_mask: Optional[torch.Tensor] = None,
1086+ **kwargs,
1087+ ) -> torch.Tensor:
1088+ """Run Gated DeltaNet on local sequence shards with all-gather state summaries."""
1089+ del kwargs
1090+ hidden_states = _local_tensor_at_cp_boundary(hidden_states)
1091+ 
1092+ base = self.module
1093+ if attention_mask is not None and attention_mask.ndim == 2:
1094+ hidden_states = hidden_states * attention_mask[:, :, None].to(
1095+ hidden_states.dtype
1096+ )
1097+ 
1098+ bsz, local_seq_len, _ = hidden_states.shape
1099+ mixed_qkv = base.in_proj_qkv(hidden_states)
1100+ z = base.in_proj_z(hidden_states).reshape(
1101+ bsz,
1102+ local_seq_len,
1103+ base.num_v_heads,
1104+ base.head_v_dim,
1105+ )
1106+ b = base.in_proj_b(hidden_states)
1107+ a = base.in_proj_a(hidden_states)
1108+ 
1109+ mixed_qkv = self._conv1d_with_halo(mixed_qkv)
1110+ query, key, value = torch.split(
1111+ mixed_qkv,
1112+ [base.key_dim, base.key_dim, base.value_dim],
1113+ dim=-1,
1114+ )
1115+ query = query.reshape(bsz, local_seq_len, base.num_k_heads, base.head_k_dim)
1116+ key = key.reshape(bsz, local_seq_len, base.num_k_heads, base.head_k_dim)
1117+ value = value.reshape(bsz, local_seq_len, base.num_v_heads, base.head_v_dim)
1118+ 
1119+ beta = b.sigmoid()
1120+ g = -base.A_log.float().exp() * F.softplus(a.float() + base.dt_bias)
1121+ 
1122+ if base.kv_groups > 1:
1123+ query = query.repeat_interleave(base.kv_groups, dim=2)
1124+ key = key.repeat_interleave(base.kv_groups, dim=2)
1125+ 
1126+ core_attn_out = _gdn_state_all_gather(
1127+ query,
1128+ key,
1129+ value,
1130+ g,
1131+ beta,
1132+ self.cp_mesh,
1133+ self.cp_rank,
1134+ self.cp_size,
1135+ use_qk_l2norm_in_kernel=True,
1136+ )
1137+ 
1138+ core_attn_out = core_attn_out.reshape(-1, base.head_v_dim)
1139+ z_flat = z.reshape(-1, base.head_v_dim)
1140+ core_attn_out = base.norm(core_attn_out, z_flat)
1141+ core_attn_out = core_attn_out.reshape(bsz, local_seq_len, base.value_dim)
1142+ if hasattr(base, "out_proj_input"):
1143+ core_attn_out = base.out_proj_input(core_attn_out)
1144+ return base.out_proj(core_attn_out)
1145+ 
1146+ 
1147+class LinearAttentionContextParallel(ParallelStyle):
1148+ """Apply context parallel execution to a Gated DeltaNet module."""
1149+ 
1150+ def __init__(self, *, mode: str = "ulysses") -> None:
1151+ if mode not in {"ulysses", "p2p", "all_gather"}:
1152+ raise NotImplementedError(
1153+ "LinearAttentionContextParallel currently supports mode='ulysses', "
1154+ "mode='p2p', and mode='all_gather'."
1155+ )
1156+ self.mode = mode
1157+ 
1158+ def apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
1159+ """Patch ``module.forward`` with a linear-attention CP executor."""
1160+ if self.mode == "ulysses":
1161+ executor = LinearAttentionUlyssesCPWrapper(module, device_mesh)
1162+ elif self.mode == "all_gather":
1163+ executor = LinearAttentionAllGatherCPWrapper(module, device_mesh)
1164+ else:
1165+ executor = LinearAttentionP2PCPWrapper(module, device_mesh)
1166+ object.__setattr__(module, "_hp_linear_attention_cp_executor", executor)
1167+ object.__setattr__(module, "_hp_linear_attention_original_forward", module.forward)
1168+ 
1169+ def _forward(*args, **kwargs):
1170+ return executor(*args, **kwargs)
1171+ 
1172+ object.__setattr__(module, "forward", _forward)
1173+ return module
Mhyper_parallel/models/qwen3_5/parallelize.py+36-68
@@ -39,6 +39,9 @@ from hyper_parallel import (
39 fully_shard,39 fully_shard,
40 parallelize_module,40 parallelize_module,
41)41)
42+from hyper_parallel.core.context_parallel.linear_attention_context_parallel import (
43+ LinearAttentionContextParallel,
44+)
42from hyper_parallel.core.pipeline_parallel import (45from hyper_parallel.core.pipeline_parallel import (
43 BatchDimSpec, Schedule1F1B, ScheduleGPipe, ScheduleInterleaved1F1B)46 BatchDimSpec, Schedule1F1B, ScheduleGPipe, ScheduleInterleaved1F1B)
44from hyper_parallel.core.pipeline_parallel.stage import SharedParameterInfo47from hyper_parallel.core.pipeline_parallel.stage import SharedParameterInfo
@@ -473,62 +476,9 @@ def qwen3_5_tp_load_transforms(
473 return transforms476 return transforms
474 477 
475 478 
476-def _redistribute_first_tensor(479+def _apply_linear_attention_cp(module: nn.Module, cp_mesh: DeviceMesh, mode: str) -> None:
477- tensor: platform.Tensor,480+ """Apply CP to a Qwen3.5 linear-attention module."""
478- mesh: DeviceMesh,481+ LinearAttentionContextParallel(mode=mode).apply(module, cp_mesh)
479- input_layout,
480- desired_layout,
481- *,
482- use_local_output: bool,
483-):
484- """Redistribute one tensor at a Qwen3.5 model-parallel hook boundary."""
485- if isinstance(tensor, DTensor):
486- dtensor = tensor
487- else:
488- dtensor = DTensor.from_local(tensor, mesh, [input_layout])
489- if tuple(dtensor.placements) != (desired_layout,):
490- dtensor = dtensor.redistribute(mesh, [desired_layout])
491- return dtensor.to_local() if use_local_output else dtensor
492- 
493- 
494-def _apply_linear_attention_cp(module: nn.Module, cp_mesh: DeviceMesh) -> None:
495- """Gather full sequence for Qwen3.5 linear attention and slice before ``out_proj``."""
496- 
497- def _pre_hook(hook_module, args, kwargs):
498- del hook_module
499- if args:
500- hidden_states = args[0]
501- rest = args[1:]
502- else:
503- hidden_states = kwargs.get("hidden_states")
504- rest = None
505- if hidden_states is None:
506- raise ValueError("linear attention CP hook expects hidden_states")
507- hidden_states = _redistribute_first_tensor(
508- hidden_states,
509- cp_mesh,
510- Shard(1),
511- Replicate(),
512- use_local_output=True,
513- )
514- if rest is None:
515- kwargs = dict(kwargs)
516- kwargs["hidden_states"] = hidden_states
517- return args, kwargs
518- return (hidden_states, *rest), kwargs
519- 
520- def _post_hook(hook_module, hook_args, output):
521- del hook_module, hook_args
522- return _redistribute_first_tensor(
523- output,
524- cp_mesh,
525- Replicate(),
526- Shard(1),
527- use_local_output=True,
528- )
529- 
530- module.register_forward_pre_hook(_pre_hook, with_kwargs=True)
531- module.out_proj_input.register_forward_hook(_post_hook)
532 482 
533 483 
534def _validate_qwen3_5_tp_config(model: Qwen3_5ForCausalLM, tp_world: int) -> None:484def _validate_qwen3_5_tp_config(model: Qwen3_5ForCausalLM, tp_world: int) -> None:
@@ -729,6 +679,7 @@ def parallelize_qwen3_5_cp(
729 cp_mesh: DeviceMesh,679 cp_mesh: DeviceMesh,
730 *,680 *,
731 ulysses_degree: Optional[int] = None,681 ulysses_degree: Optional[int] = None,
682+ linear_attention_cp_mode: str = "ulysses",
732) -> Qwen3_5ForCausalLM:683) -> Qwen3_5ForCausalLM:
733 """Apply context parallelism across the Qwen3.5 hybrid decoder.684 """Apply context parallelism across the Qwen3.5 hybrid decoder.
734 685 
@@ -746,11 +697,11 @@ def parallelize_qwen3_5_cp(
746 rank ends up with the full sequence on a head-shard and a square causal697 rank ends up with the full sequence on a head-shard and a square causal
747 mask is correct again.698 mask is correct again.
748 699 
749- Linear-attention (:class:`Qwen3_5GatedDeltaNet`) layers cannot use the700+ Linear-attention (:class:`Qwen3_5GatedDeltaNet`) layers use a matching
750- Ulysses head all-to-all (their per-head conv / SSM weights are head-fixed),701+ pure-Ulysses execution wrapper: project local sequence shards, all-to-all
751- so the model-level plan gathers the full sequence at the module boundary702+ the projected Q/K/V/B/A tensors to full-sequence local-head shards, run
752- and slices the output back to the per-rank shard before any outer output703+ the per-head conv and gated delta rule on local heads, then all-to-all the
753- layout hooks run.704+ result back to sequence shards before the output projection.
754 """705 """
755 # Only pure Ulysses is wired here; a smaller ``ulysses_degree`` makes each706 # Only pure Ulysses is wired here; a smaller ``ulysses_degree`` makes each
756 # rank attend over gathered K/V with ``is_causal=True`` but without a707 # rank attend over gathered K/V with ``is_causal=True`` but without a
@@ -773,12 +724,12 @@ def parallelize_qwen3_5_cp(
773 cp_plan.apply(block.self_attn.sdpa_core, cp_mesh)724 cp_plan.apply(block.self_attn.sdpa_core, cp_mesh)
774 full_attached += 1725 full_attached += 1
775 else:726 else:
776- _apply_linear_attention_cp(block.linear_attn, cp_mesh)727+ _apply_linear_attention_cp(block.linear_attn, cp_mesh, linear_attention_cp_mode)
777 linear_attached += 1728 linear_attached += 1
778 logger.info_rank0(729 logger.info_rank0(
779 "CP applied to Qwen3.5: cp_size=%d, ulysses_degree=%s, full-attn hooks=%d, "730 "CP applied to Qwen3.5: cp_size=%d, ulysses_degree=%s, full-attn hooks=%d, "
780- "linear-attn gather/slice=%d",731+ "linear-attn %s hooks=%d",
781- cp_mesh.size(), ulysses_degree, full_attached, linear_attached,732+ cp_mesh.size(), ulysses_degree, full_attached, linear_attention_cp_mode, linear_attached,
782 )733 )
783 return model734 return model
784 735 
@@ -786,11 +737,13 @@ def parallelize_qwen3_5_cp(
786def _needs_gqa_kv_expand_for_cp(model: Qwen3_5ForCausalLM, cp_mesh: DeviceMesh) -> bool:737def _needs_gqa_kv_expand_for_cp(model: Qwen3_5ForCausalLM, cp_mesh: DeviceMesh) -> bool:
787 """Return whether GQA K/V heads must be expanded before Ulysses CP."""738 """Return whether GQA K/V heads must be expanded before Ulysses CP."""
788 tp_size = int(getattr(model, "hp_loss_tp_scale_size", 1) or 1)739 tp_size = int(getattr(model, "hp_loss_tp_scale_size", 1) or 1)
789- if tp_size <= 1:740+ q_heads = int(model.config.num_attention_heads or 1)
790- return False
791 kv_heads = int(model.config.num_key_value_heads or 1)741 kv_heads = int(model.config.num_key_value_heads or 1)
792 local_kv_heads = max(kv_heads // tp_size, 1)742 local_kv_heads = max(kv_heads // tp_size, 1)
793- return local_kv_heads % cp_mesh.size() != 0743+ if local_kv_heads % cp_mesh.size() == 0:
744+ return False
745+ local_q_heads = max(q_heads // tp_size, 1)
746+ return local_q_heads % cp_mesh.size() == 0 and local_q_heads % local_kv_heads == 0
794 747 
795 748 
796def _resolve_fsdp_mesh(mesh):749def _resolve_fsdp_mesh(mesh):
@@ -1318,6 +1271,11 @@ def parallelize_qwen3_5(
1318 )1271 )
1319 1272 
1320 cp_size = int(cfg.train.accelerator.cp)1273 cp_size = int(cfg.train.accelerator.cp)
1274+ if tp_size > 1 and cp_size > 1 and _has_linear_attention_layers(model):
1275+ raise NotImplementedError(
1276+ "Qwen3.5 TP+CP for linear-attention layers is not supported in the "
1277+ "initial linear-attention CP path. Set parallel.tp=1 when parallel.cp>1."
1278+ )
1321 if cp_size > 1:1279 if cp_size > 1:
1322 try:1280 try:
1323 cp_mesh = mesh["cp"]1281 cp_mesh = mesh["cp"]
@@ -1329,7 +1287,17 @@ def parallelize_qwen3_5(
1329 # ``None`` (the default) resolves to pure Ulysses (degree == cp_size)1287 # ``None`` (the default) resolves to pure Ulysses (degree == cp_size)
1330 # inside ``parallelize_qwen3_5_cp``; only coerce when explicitly set.1288 # inside ``parallelize_qwen3_5_cp``; only coerce when explicitly set.
1331 ulysses_degree = int(ulysses_degree) if ulysses_degree is not None else None1289 ulysses_degree = int(ulysses_degree) if ulysses_degree is not None else None
1332- parallelize_qwen3_5_cp(model, cp_mesh, ulysses_degree=ulysses_degree)1290+ linear_attention_cp_mode = getattr(
1291+ cfg.train.accelerator,
1292+ "linear_attention_cp_mode",
1293+ "ulysses",
1294+ )
1295+ parallelize_qwen3_5_cp(
1296+ model,
1297+ cp_mesh,
1298+ ulysses_degree=ulysses_degree,
1299+ linear_attention_cp_mode=linear_attention_cp_mode,
1300+ )
1333 1301 
1334 if cfg.train.accelerator.ep > 1:1302 if cfg.train.accelerator.ep > 1:
1335 raise NotImplementedError("Qwen3.5 dense has no experts; set parallel.ep=1.")1303 raise NotImplementedError("Qwen3.5 dense has no experts; set parallel.ep=1.")
Atests/torch/context_parallel/_test_linear_attention_context_parallel.py+231-0
@@ -0,0 +1,231 @@
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+"""Distributed accuracy tests for Qwen3.5 linear-attention CP."""
16+import torch
17+import torch.distributed as dist
18+ 
19+import hyper_parallel as hp
20+from hyper_parallel.core.context_parallel.linear_attention_context_parallel import (
21+ LinearAttentionContextParallel,
22+ _differentiable_all_to_all_shard,
23+)
24+from hyper_parallel.models.qwen3_5.model import Qwen3_5GatedDeltaNet
25+from tests.torch.utils import init_dist
26+ 
27+ 
28+_MODES = ("ulysses", "p2p", "all_gather")
29+_OUTPUT_MAX_ABS_TOL = 5e-2
30+_INPUT_GRAD_MAX_ABS_TOL = 5e-1
31+_MAX_REL_SCALE_TOL = 1e-1
32+_REL_L2_TOL = 5e-2
33+_PARAM_GRAD_REL_L2_TOL = 1e-1
34+_GRAD_NORM_REL_TOL = 1e-2
35+ 
36+ 
37+def _global_max(value):
38+ value = value.detach().float()
39+ dist.all_reduce(value, op=dist.ReduceOp.MAX)
40+ return value
41+ 
42+ 
43+def _global_relative_l2(actual, expected):
44+ diff_norm_sq = (actual.detach().float() - expected.detach().float()).square().sum()
45+ expected_norm_sq = expected.detach().float().square().sum()
46+ dist.all_reduce(diff_norm_sq, op=dist.ReduceOp.SUM)
47+ dist.all_reduce(expected_norm_sq, op=dist.ReduceOp.SUM)
48+ return diff_norm_sq.sqrt() / expected_norm_sq.sqrt().clamp_min(1e-12)
49+ 
50+ 
51+def _build_module(device):
52+ return Qwen3_5GatedDeltaNet(
53+ hidden_size=128,
54+ num_v_heads=8,
55+ num_k_heads=4,
56+ head_k_dim=16,
57+ head_v_dim=16,
58+ conv_kernel_size=4,
59+ ).to(device=device, dtype=torch.bfloat16)
60+ 
61+ 
62+def _check_single_token_a2a_round_trip(mesh, rank, device):
63+ """Cover batch=1/local-sequence=1 in both Ulysses redistribution directions."""
64+ local = (
65+ torch.arange(rank * 4, (rank + 1) * 4, device=device, dtype=torch.float32)
66+ .reshape(1, 1, 4)
67+ .requires_grad_(True)
68+ )
69+ head_shard = _differentiable_all_to_all_shard(
70+ local,
71+ mesh,
72+ split_dim=2,
73+ concat_dim=1,
74+ )
75+ restored = _differentiable_all_to_all_shard(
76+ head_shard,
77+ mesh,
78+ split_dim=1,
79+ concat_dim=2,
80+ )
81+ 
82+ torch.testing.assert_close(restored, local)
83+ restored.sum().backward()
84+ torch.testing.assert_close(local.grad, torch.ones_like(local))
85+ 
86+ 
87+def _run_mode(mode, mesh, rank, world, device):
88+ """Compare one linear-attention CP mode with a full-sequence reference."""
89+ torch.manual_seed(20260726)
90+ reference = _build_module(device)
91+ candidate = _build_module(device)
92+ candidate.load_state_dict(reference.state_dict())
93+ LinearAttentionContextParallel(mode=mode).apply(candidate, mesh)
94+ 
95+ full_seq = 128
96+ local_seq = full_seq // world
97+ local_slice = slice(rank * local_seq, (rank + 1) * local_seq)
98+ torch.manual_seed(20260727)
99+ full_input = torch.randn(
100+ 1,
101+ full_seq,
102+ 128,
103+ device=device,
104+ dtype=torch.bfloat16,
105+ requires_grad=True,
106+ )
107+ local_input = (
108+ full_input[:, local_slice].detach().clone().contiguous().requires_grad_(True)
109+ )
110+ 
111+ expected_full = reference(full_input)
112+ actual = candidate(local_input)
113+ expected = expected_full[:, local_slice]
114+ output_max_abs = _global_max((actual - expected).abs().max())
115+ output_expected_max_abs = _global_max(expected.abs().max())
116+ output_max_rel_scale = output_max_abs / output_expected_max_abs.clamp_min(1e-12)
117+ output_rel_l2 = _global_relative_l2(actual, expected)
118+ 
119+ torch.manual_seed(20260800 + rank)
120+ local_grad_output = torch.randn_like(actual)
121+ gathered_grad_output = [torch.empty_like(local_grad_output) for _ in range(world)]
122+ dist.all_gather(gathered_grad_output, local_grad_output)
123+ expected_full.backward(torch.cat(gathered_grad_output, dim=1))
124+ actual.backward(local_grad_output)
125+ 
126+ input_grad_max_abs = _global_max(
127+ (local_input.grad - full_input.grad[:, local_slice]).abs().max()
128+ )
129+ input_grad_expected_max_abs = _global_max(
130+ full_input.grad[:, local_slice].abs().max()
131+ )
132+ input_grad_max_rel_scale = (
133+ input_grad_max_abs / input_grad_expected_max_abs.clamp_min(1e-12)
134+ )
135+ input_grad_rel_l2 = _global_relative_l2(
136+ local_input.grad,
137+ full_input.grad[:, local_slice],
138+ )
139+ param_grad_max_abs = torch.zeros((), device=device, dtype=torch.float32)
140+ param_grad_max_rel_scale = torch.zeros_like(param_grad_max_abs)
141+ param_grad_max_rel_l2 = torch.zeros_like(param_grad_max_abs)
142+ max_rel_scale_param = ""
143+ max_rel_l2_param = ""
144+ expected_grad_norm_sq = torch.zeros_like(param_grad_max_abs)
145+ actual_grad_norm_sq = torch.zeros_like(param_grad_max_abs)
146+ reference_params = dict(reference.named_parameters())
147+ for name, parameter in candidate.named_parameters():
148+ assert parameter.grad is not None, f"{name}.grad is missing"
149+ assert torch.isfinite(parameter.grad).all(), f"{name}.grad is not finite"
150+ actual_grad = parameter.grad.detach().float().clone()
151+ dist.all_reduce(actual_grad, op=dist.ReduceOp.SUM)
152+ assert torch.isfinite(actual_grad).all(), f"reduced {name}.grad is not finite"
153+ assert reference_params[name].grad is not None, f"reference {name}.grad is missing"
154+ expected_grad = reference_params[name].grad.detach().float()
155+ param_grad_max_abs = torch.maximum(
156+ param_grad_max_abs,
157+ (actual_grad - expected_grad).abs().max(),
158+ )
159+ grad_diff = actual_grad - expected_grad
160+ expected_max_abs = expected_grad.abs().max()
161+ expected_l2 = expected_grad.square().sum().sqrt()
162+ current_max_rel_scale = (
163+ grad_diff.abs().max() / expected_max_abs.clamp_min(1e-12)
164+ )
165+ current_rel_l2 = (
166+ grad_diff.square().sum().sqrt()
167+ / expected_l2.clamp_min(1e-12)
168+ )
169+ if current_max_rel_scale > param_grad_max_rel_scale:
170+ param_grad_max_rel_scale = current_max_rel_scale
171+ max_rel_scale_param = name
172+ if current_rel_l2 > param_grad_max_rel_l2:
173+ param_grad_max_rel_l2 = current_rel_l2
174+ max_rel_l2_param = name
175+ expected_grad_norm_sq += expected_grad.square().sum()
176+ actual_grad_norm_sq += actual_grad.square().sum()
177+ 
178+ param_grad_max_abs = _global_max(param_grad_max_abs)
179+ param_grad_max_rel_scale = _global_max(param_grad_max_rel_scale)
180+ param_grad_max_rel_l2 = _global_max(param_grad_max_rel_l2)
181+ grad_norm_rel = (
182+ (actual_grad_norm_sq.sqrt() - expected_grad_norm_sq.sqrt()).abs()
183+ / expected_grad_norm_sq.sqrt().clamp_min(1e-12)
184+ )
185+ grad_norm_rel = _global_max(grad_norm_rel)
186+ if rank == 0:
187+ print(
188+ f"mode={mode} "
189+ f"output_max_abs={output_max_abs.item():.6e} "
190+ f"output_max_rel_scale={output_max_rel_scale.item():.6e} "
191+ f"output_rel_l2={output_rel_l2.item():.6e} "
192+ f"input_grad_max_abs={input_grad_max_abs.item():.6e} "
193+ f"input_grad_max_rel_scale={input_grad_max_rel_scale.item():.6e} "
194+ f"input_grad_rel_l2={input_grad_rel_l2.item():.6e} "
195+ f"param_grad_max_abs={param_grad_max_abs.item():.6e} "
196+ f"param_grad_max_rel_scale={param_grad_max_rel_scale.item():.6e} "
197+ f"param_grad_max_rel_l2={param_grad_max_rel_l2.item():.6e} "
198+ f"grad_norm_rel={grad_norm_rel.item():.6e}",
199+ flush=True,
200+ )
201+ print(
202+ f"mode={mode} max_rel_scale_param={max_rel_scale_param} "
203+ f"max_rel_l2_param={max_rel_l2_param}",
204+ flush=True,
205+ )
206+ 
207+ # BF16 kernels may change reduction order across eager and fused backends.
208+ # Loose max-abs guards catch outliers; relative L2 guards broad numerical drift.
209+ assert output_max_abs.item() <= _OUTPUT_MAX_ABS_TOL
210+ assert input_grad_max_abs.item() <= _INPUT_GRAD_MAX_ABS_TOL
211+ assert output_max_rel_scale.item() <= _MAX_REL_SCALE_TOL
212+ assert input_grad_max_rel_scale.item() <= _MAX_REL_SCALE_TOL
213+ assert param_grad_max_rel_scale.item() <= _MAX_REL_SCALE_TOL
214+ assert output_rel_l2.item() <= _REL_L2_TOL
215+ assert input_grad_rel_l2.item() <= _REL_L2_TOL
216+ assert param_grad_max_rel_l2.item() <= _PARAM_GRAD_REL_L2_TOL
217+ assert grad_norm_rel.item() <= _GRAD_NORM_REL_TOL
218+ 
219+ 
220+def test_linear_attention_cp_forward_backward_accuracy():
221+ """All three CP modes match a full-sequence Gated DeltaNet reference."""
222+ rank, device_id = init_dist()
223+ world = dist.get_world_size()
224+ assert world == 2
225+ device = torch.device(f"npu:{device_id}")
226+ mesh = hp.init_device_mesh("npu", (world,), mesh_dim_names=("cp",))
227+ 
228+ _check_single_token_a2a_round_trip(mesh, rank, device)
229+ for mode in _MODES:
230+ _run_mode(mode, mesh, rank, world, device)
231+ dist.barrier()
Atests/torch/context_parallel/test_linear_attention_context_parallel.py+45-0
@@ -0,0 +1,45 @@
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+"""Launcher for Qwen3.5 linear-attention context parallel tests."""
16+from pathlib import Path
17+ 
18+from tests.common.mark_utils import arg_mark
19+from tests.common.parallel_case import parallel_run, TorchCase
20+ 
21+ 
22+_WORKER = str(Path(__file__).resolve().parent / "_test_linear_attention_context_parallel.py")
23+ 
24+ 
25+@arg_mark(
26+ plat_marks=["platform_ascend910b"],
27+ level_mark="level1",
28+ card_mark="allcards",
29+ essential_mark="essential",
30+)
31+def test_linear_attention_context_parallel_accuracy():
32+ """
33+ Feature: Qwen3.5 Gated DeltaNet context parallel execution
34+ Description: Compare Ulysses, P2P, and AllGather forward/backward with a full-sequence reference.
35+ Expectation: Outputs, input gradients, parameter gradients, and gradient norms match.
36+ """
37+ parallel_run(
38+ [
39+ TorchCase(
40+ _WORKER,
41+ "test_linear_attention_cp_forward_backward_accuracy",
42+ num_proc=2,
43+ )
44+ ]
45+ )
Atests/ut/core/context_parallel/test_linear_attention_context_parallel.py+83-0
@@ -0,0 +1,83 @@
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+"""Unit tests for linear-attention context parallel helpers."""
16+import pytest
17+import torch
18+ 
19+from hyper_parallel.core.context_parallel.linear_attention_context_parallel import (
20+ LinearAttentionContextParallel,
21+ _merge_gdn_prefix_state_summaries_torch,
22+ _pack_gdn_state_summary,
23+ _slice_qkv_local_cp,
24+ _unpack_gdn_state_summary,
25+)
26+ 
27+ 
28+@pytest.mark.parametrize("mode", ("ulysses", "p2p", "all_gather"))
29+def test_linear_attention_cp_accepts_supported_modes(mode):
30+ """All public linear-attention CP modes can be constructed."""
31+ assert LinearAttentionContextParallel(mode=mode).mode == mode
32+ 
33+ 
34+def test_linear_attention_cp_rejects_unknown_mode():
35+ """An unsupported execution mode fails before patching a module."""
36+ with pytest.raises(NotImplementedError, match="currently supports"):
37+ LinearAttentionContextParallel(mode="unknown")
38+ 
39+ 
40+def test_slice_qkv_local_cp_slices_each_projection_independently():
41+ """Fused Q/K/V projections preserve their channel boundaries when sharded."""
42+ tensor = torch.arange(20, dtype=torch.float32).reshape(1, 1, 20)
43+ 
44+ actual = _slice_qkv_local_cp(
45+ tensor,
46+ key_dim=6,
47+ value_dim=8,
48+ dim=-1,
49+ cp_rank=1,
50+ cp_size=2,
51+ )
52+ 
53+ expected = torch.tensor([[[3, 4, 5, 9, 10, 11, 16, 17, 18, 19]]], dtype=torch.float32)
54+ torch.testing.assert_close(actual, expected)
55+ 
56+ 
57+def test_gdn_summary_pack_merge_and_backward():
58+ """Packed summaries compose as affine state transitions and remain differentiable."""
59+ state_ext = torch.tensor(
60+ [[[[[1.0], [2.0]]]], [[[[3.0], [4.0]]]]],
61+ requires_grad=True,
62+ )
63+ transition = torch.tensor(
64+ [
65+ [[[[2.0, 0.0], [0.0, 3.0]]]],
66+ [[[[4.0, 0.0], [0.0, 5.0]]]],
67+ ],
68+ requires_grad=True,
69+ )
70+ 
71+ packed = _pack_gdn_state_summary(state_ext, transition)
72+ unpacked_state, unpacked_transition = _unpack_gdn_state_summary(packed, v_head_dim=1)
73+ actual = _merge_gdn_prefix_state_summaries_torch(
74+ unpacked_state,
75+ unpacked_transition,
76+ rank=2,
77+ )
78+ 
79+ expected = transition[1] @ state_ext[0] + state_ext[1]
80+ torch.testing.assert_close(actual, expected)
81+ actual.sum().backward()
82+ assert state_ext.grad is not None
83+ assert transition.grad is not None