« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Context parallel execution for Qwen3.5-style Gated DeltaNet layers."""
16from __future__ import annotations
18from typing import NamedTuple, Optional
20import torch
21import torch.distributed as dist
22from torch import nn
23from torch.nn import functional as F
24from torch.utils.checkpoint import checkpoint
26from hyper_parallel.core.context_parallel.context_parallel import (
27 _ensure_1d,
28)
29from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
30from hyper_parallel.core.dtensor.dtensor import DTensor
31from hyper_parallel.core.tensor_parallel.style import ParallelStyle
32from hyper_parallel.models.modules.linear_attention import torch_chunk_gated_delta_rule
33from hyper_parallel.platform import get_platform
36platform = get_platform()
39def _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])
44def _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)
61def _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 )
82def _local_tensor_at_cp_boundary(tensor: torch.Tensor) -> torch.Tensor:
83 """Return the local tensor carried by a CP-boundary input.
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
95def _all_to_all_previous_rank_halo(
96 tail: torch.Tensor,
97 cp_mesh: DeviceMesh,
98 cp_rank: int,
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)
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]
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]
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
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()
133def _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)
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 )
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)
173def _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 )
190def _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)
195class _GDNPreparedChunks(NamedTuple):
196 """Reusable chunk intermediates shared by state-summary CP modes."""
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
210def _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)
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 ]
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
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)
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)
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 )
278def _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()
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]
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
314 return state_ext, transition
317def _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)
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)
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 )
350def _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)
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 )
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)
406def _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)
419def _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
434def _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}.")
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
456def _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
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 )
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)
508class _RecvInitialStateP2PFunction(torch.autograd.Function):
509 """Receive the recurrent initial state; send its gradient in backward."""
511 @staticmethod
512 def forward( # pylint: disable=arguments-differ
513 ctx,
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
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
534class _SendFinalStateP2PFunction(torch.autograd.Function):
535 """Send the recurrent final state; receive its gradient in backward."""
537 @staticmethod
538 def forward( # pylint: disable=arguments-differ
539 ctx,
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(())
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
563def _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
574def _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.
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
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])
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)
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 )
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)
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
646 return core_attn_out
649def _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``.
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
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 )
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)
682 reshape_dims = list(original_shape)
683 reshape_dims[split_dim] = split_count
684 reshape_dims.insert(split_dim + 1, split_size)
686 trans_dims = list(range(len(reshape_dims)))
687 trans_dims.remove(split_dim)
688 trans_dims.insert(0, split_dim)
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)
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 )
707 output_reshape = list(output.shape)
708 output_reshape[0] = split_count
709 output_reshape.insert(1, output.shape[0] // split_count)
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)
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)
726 return final_output.reshape(final_reshape).view(final_shape)
729class LinearAttentionUlyssesCPWrapper(nn.Module):
730 """Pure-Ulysses CP execution wrapper for a Qwen3.5 Gated DeltaNet module.
732 Parameters stay owned by the original module. The wrapper only changes the
733 execution layout:
735 ``[B, S_local, full_heads] -> [B, S_full, local_heads] ->
736 [B, S_local, full_heads]``.
737 """
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()
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 )
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 )
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 )
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
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
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 )
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 )
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 )
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)
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 )
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)
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)
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
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)
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)
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)
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)
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 )
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)
934class LinearAttentionP2PCPWrapper(nn.Module):
935 """Sequence-sharded GDN CP with an affine-summary state wavefront."""
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()
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 )
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 )
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)
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 )
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)
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)
1011 beta = b.sigmoid()
1012 g = -base.A_log.float().exp() * F.softplus(a.float() + base.dt_bias)
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)
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 )
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)
1039class LinearAttentionAllGatherCPWrapper(nn.Module):
1040 """Sequence-sharded GDN CP using all-gathered recurrent-state summaries."""
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()
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 )
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 )
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)
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 )
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)
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)
1119 beta = b.sigmoid()
1120 g = -base.A_log.float().exp() * F.softplus(a.float() + base.dt_bias)
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)
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 )
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)
1147class LinearAttentionContextParallel(ParallelStyle):
1148 """Apply context parallel execution to a Gated DeltaNet module."""
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
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)
1169 def _forward(*args, **kwargs):
1170 return executor(*args, **kwargs)
1172 object.__setattr__(module, "forward", _forward)
1173 return module