« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +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."""
17# This module is the Torch implementation of the public CP style.
18# pylint: disable=forbidden-backend-import,missing-public-type-hints
19# pylint: disable=missing-public-docstring,not-callable
20# PyTorch autograd.Function intentionally defines framework-specific signatures.
21# pylint: disable=abstract-method,arguments-differ
22from __future__ import annotations
24from typing import NamedTuple, Optional
26import torch
27import torch.distributed as dist
28from torch import nn
29from torch.nn import functional as F
30from torch.utils.checkpoint import checkpoint
32from hyper_parallel.core.context_parallel.context_parallel import (
33 _ensure_1d,
34)
35from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
36from hyper_parallel.core.dtensor.dtensor import DTensor
37from hyper_parallel.core.tensor_parallel.style import ParallelStyle
38from hyper_parallel.models.modules.linear_attention import (
39 chunk_gated_delta_rule,
40 is_triton_gdn_available,
41 torch_chunk_gated_delta_rule,
42)
43from hyper_parallel.platform import get_platform
46platform = get_platform()
49def _global_peer_rank(cp_mesh: DeviceMesh, local_rank: int) -> int:
50 """Map a CP-local rank index to its global distributed rank."""
51 return int(cp_mesh.rank_list[local_rank])
54def _slice_local_cp(
55 tensor: torch.Tensor,
56 dim: int,
57 cp_rank: int,
58 cp_size: int,
59) -> torch.Tensor:
60 """Return this CP rank's contiguous slice along ``dim``."""
61 dim_size = tensor.shape[dim]
62 if dim_size % cp_size != 0:
63 raise ValueError(
64 f"linear attention CP expects dim size {dim_size} "
65 f"to be divisible by cp_size {cp_size}."
66 )
67 chunk = dim_size // cp_size
68 return tensor.narrow(dim, cp_rank * chunk, chunk)
71def _slice_qkv_local_cp(
72 tensor: torch.Tensor,
73 *,
74 key_dim: int,
75 value_dim: int,
76 dim: int,
77 cp_rank: int,
78 cp_size: int,
79) -> torch.Tensor:
80 """Slice a fused ``[Q, K, V]`` tensor on the Q/K/V channel dimension."""
81 q, k, v = torch.split(tensor, [key_dim, key_dim, value_dim], dim=dim)
82 return torch.cat(
83 (
84 _slice_local_cp(q, dim, cp_rank, cp_size),
85 _slice_local_cp(k, dim, cp_rank, cp_size),
86 _slice_local_cp(v, dim, cp_rank, cp_size),
87 ),
88 dim=dim,
89 )
92def _local_tensor_at_cp_boundary(tensor: torch.Tensor) -> torch.Tensor:
93 """Return the local tensor carried by a CP-boundary input.
95 The first supported Qwen3.5 linear-attention CP path keeps decoder-layer
96 activations as local sequence shards. If an upstream wrapper passes that
97 shard as a DTensor, use its local tensor and continue with the same
98 ``[B, S_local, H]`` boundary contract.
99 """
100 if isinstance(tensor, DTensor):
101 return tensor.to_local()
102 return tensor
105def _all_to_all_previous_rank_halo(
106 tail: torch.Tensor,
107 cp_mesh: DeviceMesh,
108 cp_rank: int,
109 cp_size: int,
110) -> torch.Tensor:
111 """Send a convolution halo only to the next rank using differentiable A2AV."""
112 if cp_size == 1:
113 return torch.zeros_like(tail)
115 cp_group = cp_mesh.get_group()
116 group_ranks = tuple(int(rank) for rank in dist.get_process_group_ranks(cp_group))
117 rank_list = tuple(int(rank) for rank in cp_mesh.rank_list)
118 rank_to_group_index = {rank: index for index, rank in enumerate(group_ranks)}
119 halo_width = tail.shape[1]
121 input_splits = [0] * cp_size
122 exchange_input = tail.permute(1, 0, 2).contiguous()
123 if cp_rank < cp_size - 1:
124 input_splits[rank_to_group_index[rank_list[cp_rank + 1]]] = halo_width
125 else:
126 exchange_input = exchange_input[:0]
128 output_splits = [0] * cp_size
129 if cp_rank > 0:
130 output_splits[rank_to_group_index[rank_list[cp_rank - 1]]] = halo_width
132 exchange_output = platform.differentiable_all_to_all_single(
133 exchange_input,
134 input_splits,
135 output_splits,
136 group=cp_group,
137 )
138 if cp_rank == 0:
139 return torch.zeros_like(tail) + exchange_output.sum().to(tail.dtype) * 0
140 return exchange_output.permute(1, 0, 2).contiguous()
143def _causal_conv1d_with_cp_halo(
144 mixed_qkv: torch.Tensor,
145 conv1d: nn.Conv1d,
146 cp_mesh: DeviceMesh,
147 cp_rank: int,
148 cp_size: int,
149) -> torch.Tensor:
150 """Run causal depthwise Conv1d with only the previous rank's boundary."""
151 kernel_size = conv1d.kernel_size[0]
152 dilation = conv1d.dilation[0]
153 halo_width = (kernel_size - 1) * dilation
154 if halo_width == 0 or cp_size == 1:
155 conv_out = conv1d(mixed_qkv.transpose(1, 2))
156 return F.silu(conv_out[:, :, : mixed_qkv.shape[1]]).transpose(1, 2)
158 if mixed_qkv.shape[1] < halo_width:
159 raise ValueError(
160 "linear attention CP conv halo requires local_seq_len >= "
161 f"{halo_width}, got {mixed_qkv.shape[1]}."
162 )
164 halo = _all_to_all_previous_rank_halo(
165 mixed_qkv[:, -halo_width:, :].contiguous(),
166 cp_mesh,
167 cp_rank,
168 cp_size,
169 )
170 conv_input = torch.cat((halo, mixed_qkv), dim=1).transpose(1, 2)
171 conv_out = F.conv1d(
172 input=conv_input,
173 weight=conv1d.weight,
174 bias=conv1d.bias,
175 stride=conv1d.stride,
176 padding=0,
177 dilation=conv1d.dilation,
178 groups=conv1d.groups,
179 )
180 return F.silu(conv_out).transpose(1, 2)
183def _all_gather_stack(
184 tensor: torch.Tensor,
185 cp_mesh: DeviceMesh,
186 cp_size: int,
187) -> torch.Tensor:
188 """All-gather equal-shaped tensors and stack them on a leading rank dim."""
189 if cp_size == 1:
190 return tensor.unsqueeze(0)
191 return platform.differentiable_all_gather_concat(
192 tensor.unsqueeze(0),
193 cp_mesh.get_group(),
194 cp_size,
195 0,
196 tuple(int(rank) for rank in cp_mesh.rank_list),
197 )
200def _l2norm_torch(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor:
201 """Match the pure torch GDN reference l2norm helper."""
202 return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
205class _GDNPreparedChunks(NamedTuple):
206 """Reusable chunk intermediates shared by state-summary CP modes."""
208 initial_dtype: torch.dtype
209 query: torch.Tensor
210 key: torch.Tensor
211 chunk_value: torch.Tensor
212 g: torch.Tensor
213 decay_mask: torch.Tensor
214 k_cumdecay: torch.Tensor
215 sequence_length: int
216 total_sequence_length: int
217 chunk_size: int
220def _prepare_gdn_chunks_for_summary(
221 query: torch.Tensor,
222 key: torch.Tensor,
223 value: torch.Tensor,
224 g: torch.Tensor,
225 beta: torch.Tensor,
226 *,
227 chunk_size: int = 64,
228 use_qk_l2norm_in_kernel: bool = False,
229) -> _GDNPreparedChunks:
230 """Prepare GDN chunk intermediates shared by summary and local output."""
231 initial_dtype = query.dtype
232 if use_qk_l2norm_in_kernel:
233 query = _l2norm_torch(query, dim=-1, eps=1e-6)
234 key = _l2norm_torch(key, dim=-1, eps=1e-6)
236 query, key, value, beta, g = [
237 x.transpose(1, 2).contiguous().to(torch.float32)
238 for x in (query, key, value, beta, g)
239 ]
241 sequence_length = key.shape[2]
242 pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
243 query = F.pad(query, (0, 0, 0, pad_size))
244 key = F.pad(key, (0, 0, 0, pad_size))
245 value = F.pad(value, (0, 0, 0, pad_size))
246 beta = F.pad(beta, (0, pad_size))
247 g = F.pad(g, (0, pad_size))
248 total_sequence_length = sequence_length + pad_size
250 query = query * (1 / (query.shape[-1] ** 0.5))
251 v_beta = value * beta.unsqueeze(-1)
252 k_beta = key * beta.unsqueeze(-1)
253 query, key, k_beta, v_beta = [
254 x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
255 for x in (query, key, k_beta, v_beta)
256 ]
257 g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
259 mask = torch.triu(
260 torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
261 diagonal=0,
262 )
263 g = g.cumsum(dim=-1)
264 decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
265 attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
266 for row_idx in range(1, chunk_size):
267 row = attn[..., row_idx, :row_idx].clone()
268 sub = attn[..., :row_idx, :row_idx].clone()
269 attn[..., row_idx, :row_idx] = row + (row.unsqueeze(-1) * sub).sum(-2)
270 attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
272 chunk_value = attn @ v_beta
273 k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
274 return _GDNPreparedChunks(
275 initial_dtype=initial_dtype,
276 query=query,
277 key=key,
278 chunk_value=chunk_value,
279 g=g,
280 decay_mask=decay_mask,
281 k_cumdecay=k_cumdecay,
282 sequence_length=sequence_length,
283 total_sequence_length=total_sequence_length,
284 chunk_size=chunk_size,
285 )
288def _compute_gdn_state_summary_from_prepared(
289 prepared: _GDNPreparedChunks,
290) -> tuple[torch.Tensor, torch.Tensor]:
291 """Compute ``state_out = M @ state_in + S`` from prepared GDN chunks."""
292 key = prepared.key
293 batch_size, num_heads, _, _, k_head_dim = key.shape
294 v_head_dim = prepared.chunk_value.shape[-1]
295 eye = torch.eye(k_head_dim, device=key.device, dtype=torch.float32).reshape(
296 1, 1, k_head_dim, k_head_dim
297 )
298 state_ext = torch.zeros(
299 batch_size,
300 num_heads,
301 k_head_dim,
302 v_head_dim,
303 device=key.device,
304 dtype=torch.float32,
305 )
306 transition = eye.expand(batch_size, num_heads, -1, -1).clone()
308 for chunk_idx in range(key.shape[2]):
309 key_i = key[:, :, chunk_idx]
310 value_i = prepared.chunk_value[:, :, chunk_idx]
311 w_i = prepared.k_cumdecay[:, :, chunk_idx]
312 g_i = prepared.g[:, :, chunk_idx]
313 decay = g_i[:, :, -1].exp()
314 key_decay = key_i * (g_i[:, :, -1, None] - g_i).exp()[..., None]
316 transition_i = (
317 decay[:, :, None, None] * eye
318 - key_decay.transpose(-1, -2) @ w_i
319 )
320 state_ext_i = key_decay.transpose(-1, -2) @ value_i
321 state_ext = transition_i @ state_ext + state_ext_i
322 transition = transition_i @ transition
324 return state_ext, transition
327def _checkpoint_gdn_state_summary(
328 prepared: _GDNPreparedChunks,
329) -> tuple[torch.Tensor, torch.Tensor]:
330 """Compute a state summary without retaining its per-chunk autograd graph."""
331 if not torch.is_grad_enabled():
332 return _compute_gdn_state_summary_from_prepared(prepared)
334 def recompute(
335 key: torch.Tensor,
336 chunk_value: torch.Tensor,
337 g: torch.Tensor,
338 k_cumdecay: torch.Tensor,
339 ) -> tuple[torch.Tensor, torch.Tensor]:
340 """Rebuild a prepared view from explicit checkpoint inputs."""
341 checkpoint_prepared = prepared._replace(
342 key=key,
343 chunk_value=chunk_value,
344 g=g,
345 k_cumdecay=k_cumdecay,
346 )
347 return _compute_gdn_state_summary_from_prepared(checkpoint_prepared)
349 return checkpoint(
350 recompute,
351 prepared.key,
352 prepared.chunk_value,
353 prepared.g,
354 prepared.k_cumdecay,
355 use_reentrant=False,
356 preserve_rng_state=False,
357 )
360def _run_prepared_gdn_chunks(
361 prepared: _GDNPreparedChunks,
362 initial_state: Optional[torch.Tensor],
363) -> torch.Tensor:
364 """Run local GDN output using already prepared chunk intermediates."""
365 query = prepared.query
366 key = prepared.key
367 chunk_value = prepared.chunk_value
368 batch_size, num_heads, _, _, k_head_dim = key.shape
369 v_head_dim = chunk_value.shape[-1]
370 recurrent_state = (
371 torch.zeros(
372 batch_size,
373 num_heads,
374 k_head_dim,
375 v_head_dim,
376 device=chunk_value.device,
377 dtype=chunk_value.dtype,
378 )
379 if initial_state is None
380 else initial_state.to(chunk_value)
381 )
382 core_attn_out = torch.zeros_like(chunk_value)
384 for chunk_idx in range(0, prepared.total_sequence_length // prepared.chunk_size):
385 q_i = query[:, :, chunk_idx]
386 k_i = key[:, :, chunk_idx]
387 v_i = chunk_value[:, :, chunk_idx]
388 attn = q_i @ k_i.transpose(-1, -2) * prepared.decay_mask[:, :, chunk_idx]
389 v_prime = prepared.k_cumdecay[:, :, chunk_idx] @ recurrent_state
390 v_new = v_i - v_prime
391 attn_inter = (
392 q_i * prepared.g[:, :, chunk_idx, :, None].exp()
393 ) @ recurrent_state
394 core_attn_out[:, :, chunk_idx] = attn_inter + attn @ v_new
395 recurrent_state = (
396 recurrent_state * prepared.g[:, :, chunk_idx, -1, None, None].exp()
397 + (
398 k_i
399 * (
400 prepared.g[:, :, chunk_idx, -1, None]
401 - prepared.g[:, :, chunk_idx]
402 ).exp()[..., None]
403 ).transpose(-1, -2) @ v_new
404 )
406 core_attn_out = core_attn_out.reshape(
407 core_attn_out.shape[0],
408 core_attn_out.shape[1],
409 -1,
410 core_attn_out.shape[-1],
411 )
412 core_attn_out = core_attn_out[:, :, :prepared.sequence_length]
413 return core_attn_out.transpose(1, 2).contiguous().to(prepared.initial_dtype)
416def _pack_gdn_state_summary(
417 state_ext: torch.Tensor,
418 transition: torch.Tensor,
419) -> torch.Tensor:
420 """Pack ``S`` and ``M`` summaries into one all-gather payload."""
421 if state_ext.shape[:-1] != transition.shape[:-1]:
422 raise ValueError(
423 "state_ext and transition must share [B,H,K] dimensions, got "
424 f"{tuple(state_ext.shape)} and {tuple(transition.shape)}."
425 )
426 return torch.cat((state_ext, transition), dim=-1)
429def _unpack_gdn_state_summary(
430 packed: torch.Tensor,
431 v_head_dim: int,
432) -> tuple[torch.Tensor, torch.Tensor]:
433 """Unpack a gathered ``[S, M]`` payload."""
434 if packed.shape[-1] <= v_head_dim:
435 raise ValueError(
436 f"packed state summary last dim must be > v_head_dim={v_head_dim}, "
437 f"got {packed.shape[-1]}."
438 )
439 state_ext = packed[..., :v_head_dim]
440 transition = packed[..., v_head_dim:]
441 return state_ext, transition
444def _merge_gdn_prefix_state_summaries_torch(
445 state_ext: torch.Tensor,
446 transition: torch.Tensor,
447 rank: int,
448) -> torch.Tensor:
449 """Merge gathered GDN summaries before ``rank`` into its initial state."""
450 if state_ext.dim() != 5 or transition.dim() != 5:
451 raise ValueError(
452 "state summary merge expects state_ext [R,B,H,K,V] and "
453 "transition [R,B,H,K,K]."
454 )
455 if state_ext.shape[0] != transition.shape[0]:
456 raise ValueError("state_ext and transition must have the same rank dimension.")
457 if rank < 0 or rank > state_ext.shape[0]:
458 raise ValueError(f"rank must be in [0, {state_ext.shape[0]}], got {rank}.")
460 state = torch.zeros_like(state_ext[0])
461 for prev_rank in range(rank):
462 state = transition[prev_rank] @ state + state_ext[prev_rank]
463 return state
466def _gdn_state_all_gather(
467 query: torch.Tensor,
468 key: torch.Tensor,
469 value: torch.Tensor,
470 g: torch.Tensor,
471 beta: torch.Tensor,
472 cp_mesh: DeviceMesh,
473 cp_rank: int,
474 cp_size: int,
475 *,
476 use_qk_l2norm_in_kernel: bool,
477) -> torch.Tensor:
478 """Apply local GDN with all-gathered recurrent-state summaries."""
479 if cp_size == 1:
480 core_attn_out, _ = torch_chunk_gated_delta_rule(
481 query,
482 key,
483 value,
484 g=g,
485 beta=beta,
486 initial_state=None,
487 output_final_state=False,
488 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
489 )
490 return core_attn_out
492 prepared = _prepare_gdn_chunks_for_summary(
493 query,
494 key,
495 value,
496 g,
497 beta,
498 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
499 )
500 state_ext, transition = _checkpoint_gdn_state_summary(prepared)
501 packed_summary = _pack_gdn_state_summary(state_ext, transition)
502 gathered_summary = _all_gather_stack(packed_summary, cp_mesh, cp_size)
503 gathered_state_ext, gathered_transition = _unpack_gdn_state_summary(
504 gathered_summary,
505 state_ext.shape[-1],
506 )
508 initial_state = _merge_gdn_prefix_state_summaries_torch(
509 gathered_state_ext,
510 gathered_transition,
511 cp_rank,
512 )
513 all_gather_tie = gathered_summary.sum()
514 initial_state = initial_state + all_gather_tie.to(initial_state.dtype) * 0
515 return _run_prepared_gdn_chunks(prepared, initial_state)
518class _RecvInitialStateP2PFunction(torch.autograd.Function):
519 """Receive the recurrent initial state; send its gradient in backward."""
521 @staticmethod
522 def forward( # pylint: disable=arguments-differ
523 ctx,
524 anchor: torch.Tensor,
525 cp_group,
526 prev_rank: int,
527 state_shape: tuple[int, ...],
528 ) -> torch.Tensor:
529 """Receive the initial state from the preceding CP rank."""
530 state = torch.empty(state_shape, device=anchor.device, dtype=torch.float32)
531 dist.recv(state, src=prev_rank, group=cp_group)
532 ctx.cp_group = cp_group
533 ctx.prev_rank = prev_rank
534 return state
536 @staticmethod
537 def backward(ctx, grad_state: Optional[torch.Tensor]):
538 if grad_state is None:
539 raise RuntimeError("linear attention P2P backward missing initial-state grad.")
540 dist.send(grad_state.contiguous(), dst=ctx.prev_rank, group=ctx.cp_group)
541 return None, None, None, None
544class _SendFinalStateP2PFunction(torch.autograd.Function):
545 """Send the recurrent final state; receive its gradient in backward."""
547 @staticmethod
548 def forward( # pylint: disable=arguments-differ
549 ctx,
550 final_state: torch.Tensor,
551 cp_group,
552 next_rank: int,
553 ) -> torch.Tensor:
554 """Send the final state to the succeeding CP rank."""
555 dist.send(final_state.contiguous(), dst=next_rank, group=cp_group)
556 ctx.cp_group = cp_group
557 ctx.next_rank = next_rank
558 ctx.state_shape = tuple(final_state.shape)
559 ctx.state_dtype = final_state.dtype
560 return final_state.new_zeros(())
562 @staticmethod
563 def backward(ctx, grad_token: torch.Tensor):
564 grad_state = torch.empty(
565 ctx.state_shape,
566 device=grad_token.device,
567 dtype=ctx.state_dtype,
568 )
569 dist.recv(grad_state, src=ctx.next_rank, group=ctx.cp_group)
570 return grad_state, None, None
573def _apply_gdn_state_summary(
574 state_ext: torch.Tensor,
575 transition: torch.Tensor,
576 initial_state: Optional[torch.Tensor],
577) -> torch.Tensor:
578 """Apply ``state_out = M @ state_in + S`` to an incoming GDN state."""
579 if initial_state is None:
580 return state_ext
581 return transition @ initial_state.to(transition) + state_ext
584def _gdn_state_p2p_summary(
585 query: torch.Tensor,
586 key: torch.Tensor,
587 value: torch.Tensor,
588 g: torch.Tensor,
589 beta: torch.Tensor,
590 cp_mesh: DeviceMesh,
591 cp_rank: int,
592 cp_size: int,
593 *,
594 use_qk_l2norm_in_kernel: bool,
595) -> torch.Tensor:
596 """Run local GDN with an affine-summary state wavefront.
598 Every rank prepares its local chunks and state transition in parallel.
599 The rank-ordered critical path then contains only ``M @ state + S`` and
600 the small state transfer. Token outputs retain the ordinary PyTorch graph,
601 while the two custom autograd boundaries reverse the state communication.
602 """
603 if cp_size == 1:
604 core_attn_out, _ = torch_chunk_gated_delta_rule(
605 query,
606 key,
607 value,
608 g=g,
609 beta=beta,
610 initial_state=None,
611 output_final_state=False,
612 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
613 )
614 return core_attn_out
616 cp_group = cp_mesh.get_group()
617 prev_rank = _global_peer_rank(cp_mesh, cp_rank - 1) if cp_rank > 0 else -1
618 next_rank = _global_peer_rank(cp_mesh, cp_rank + 1) if cp_rank < cp_size - 1 else -1
619 state_shape = (query.shape[0], value.shape[2], query.shape[3], value.shape[3])
621 prepared = _prepare_gdn_chunks_for_summary(
622 query,
623 key,
624 value,
625 g,
626 beta,
627 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
628 )
629 state_ext = None
630 transition = None
631 if cp_rank < cp_size - 1:
632 state_ext, transition = _checkpoint_gdn_state_summary(prepared)
634 initial_state = None
635 if cp_rank > 0:
636 initial_state = _RecvInitialStateP2PFunction.apply(
637 query,
638 cp_group,
639 prev_rank,
640 state_shape,
641 )
643 send_token = None
644 if cp_rank < cp_size - 1:
645 final_state = _apply_gdn_state_summary(
646 state_ext,
647 transition,
648 initial_state,
649 )
650 send_token = _SendFinalStateP2PFunction.apply(final_state, cp_group, next_rank)
652 core_attn_out = _run_prepared_gdn_chunks(prepared, initial_state)
653 if send_token is not None:
654 core_attn_out = core_attn_out + send_token.to(core_attn_out.dtype) * 0
656 return core_attn_out
659class _GDNStateP2PTritonFunction(torch.autograd.Function):
660 """Pipeline fused affine GDN states over sequence-sharded CP ranks."""
662 @staticmethod
663 def forward( # pylint: disable=arguments-differ,too-many-locals
664 ctx,
665 query: torch.Tensor,
666 key: torch.Tensor,
667 value: torch.Tensor,
668 g: torch.Tensor,
669 beta: torch.Tensor,
670 cp_rank: int,
671 cp_size: int,
672 cp_group,
673 prev_rank: int,
674 next_rank: int,
675 ) -> torch.Tensor:
676 """Run fused local GDN and forward its affine state across CP ranks."""
677 from hyper_parallel.platform.torch.custom_ops.gdn.chunk_gated_delta_rule import ( # pylint: disable=import-outside-toplevel
678 chunk_gated_delta_rule_fwd_apply_state_saved,
679 chunk_gated_delta_rule_fwd_output_saved,
680 chunk_gated_delta_rule_fwd_prepare_saved,
681 )
682 from hyper_parallel.platform.torch.custom_ops.gdn.state_summary import ( # pylint: disable=import-outside-toplevel
683 apply_gdn_state_summary,
684 chunk_gated_delta_rule_state_summary_fwd,
685 )
687 (
688 query_norm,
689 key_norm,
690 _,
691 _,
692 g_cumsum,
693 matrix_a,
694 w,
695 u,
696 scale,
697 ) = chunk_gated_delta_rule_fwd_prepare_saved(
698 query,
699 key,
700 value,
701 g,
702 beta,
703 use_qk_l2norm_in_kernel=False,
704 )
705 initial_state = None
706 recv_buffer = None
707 recv_work = None
708 if cp_rank > 0:
709 recv_buffer = torch.empty(
710 (query.shape[0], query.shape[2], query.shape[3], value.shape[3]),
711 device=query.device,
712 dtype=torch.float32,
713 )
714 recv_work = dist.irecv(recv_buffer, src=prev_rank, group=cp_group)
716 state_ext = None
717 transition = None
718 if cp_rank < cp_size - 1:
719 state_ext, transition = chunk_gated_delta_rule_state_summary_fwd(
720 key_norm,
721 w,
722 u,
723 g_cumsum,
724 )
726 if recv_work is not None:
727 recv_work.wait()
728 initial_state = recv_buffer
730 send_buffer = None
731 send_work = None
732 if cp_rank < cp_size - 1:
733 send_buffer = apply_gdn_state_summary(
734 state_ext,
735 transition,
736 initial_state,
737 ).contiguous()
738 send_work = dist.isend(send_buffer, dst=next_rank, group=cp_group)
740 h, v_new, _ = chunk_gated_delta_rule_fwd_apply_state_saved(
741 key_norm,
742 g_cumsum,
743 w,
744 u,
745 initial_state=initial_state,
746 output_final_state=False,
747 )
748 output = chunk_gated_delta_rule_fwd_output_saved(
749 query_norm,
750 key_norm,
751 g_cumsum,
752 h,
753 v_new,
754 scale,
755 ).to(query.dtype)
757 if send_work is not None:
758 send_work.wait()
760 empty = query.new_empty(0)
761 ctx.save_for_backward(
762 query_norm,
763 key_norm,
764 value,
765 g_cumsum,
766 beta,
767 matrix_a,
768 initial_state if initial_state is not None else empty,
769 transition if transition is not None else empty,
770 )
771 ctx.has_initial_state = initial_state is not None
772 ctx.cp_rank = cp_rank
773 ctx.cp_size = cp_size
774 ctx.cp_group = cp_group
775 ctx.prev_rank = prev_rank
776 ctx.next_rank = next_rank
777 ctx.scale = scale
778 return output
780 @staticmethod
781 def backward(ctx, grad_output: torch.Tensor): # pylint: disable=too-many-locals
782 """Backpropagate local GDN tensors and the state gradient wavefront."""
783 from hyper_parallel.platform.torch.custom_ops.gdn.chunk_gated_delta_rule import ( # pylint: disable=import-outside-toplevel
784 chunk_gated_delta_rule_bwd_finish_saved,
785 chunk_gated_delta_rule_bwd_prepare_saved,
786 chunk_gated_delta_rule_bwd_state_saved,
787 )
788 from hyper_parallel.platform.torch.custom_ops.gdn.state_summary import ( # pylint: disable=import-outside-toplevel
789 apply_gdn_state_gradient_summary,
790 chunk_gated_delta_rule_state_gradient_summary_bwd,
791 )
793 (
794 query,
795 key,
796 value,
797 g_cumsum,
798 beta,
799 matrix_a,
800 initial_state,
801 transition,
802 ) = ctx.saved_tensors
803 if not ctx.has_initial_state:
804 initial_state = None
806 w, h, v_new, dv = chunk_gated_delta_rule_bwd_prepare_saved(
807 query,
808 key,
809 value,
810 g_cumsum,
811 beta,
812 matrix_a,
813 initial_state,
814 grad_output,
815 ctx.scale,
816 )
818 grad_state_ext = None
819 if ctx.cp_rank > 0:
820 grad_state_ext = chunk_gated_delta_rule_state_gradient_summary_bwd(
821 query,
822 key,
823 w,
824 g_cumsum,
825 grad_output,
826 dv,
827 ctx.scale,
828 )
830 grad_final_state = None
831 recv_work = None
832 if ctx.cp_rank < ctx.cp_size - 1:
833 recv_buffer = torch.empty(
834 (query.shape[0], query.shape[2], query.shape[3], value.shape[3]),
835 device=grad_output.device,
836 dtype=torch.float32,
837 )
838 recv_work = dist.irecv(
839 recv_buffer,
840 src=ctx.next_rank,
841 group=ctx.cp_group,
842 )
843 if recv_work is not None:
844 recv_work.wait()
845 grad_final_state = recv_buffer
847 send_buffer = None
848 send_work = None
849 if ctx.cp_rank > 0:
850 send_buffer = apply_gdn_state_gradient_summary(
851 grad_state_ext,
852 transition,
853 grad_final_state,
854 ).contiguous()
855 send_work = dist.isend(
856 send_buffer,
857 dst=ctx.prev_rank,
858 group=ctx.cp_group,
859 )
861 dh, _, dv = chunk_gated_delta_rule_bwd_state_saved(
862 query,
863 key,
864 g_cumsum,
865 w,
866 initial_state,
867 grad_final_state,
868 grad_output,
869 dv,
870 ctx.scale,
871 )
872 empty = query.new_empty(0)
873 dq, dk, dv, dg, dbeta = chunk_gated_delta_rule_bwd_finish_saved(
874 query,
875 key,
876 query,
877 key,
878 value,
879 g_cumsum,
880 beta,
881 matrix_a,
882 w,
883 h,
884 v_new,
885 dv,
886 grad_output,
887 dh,
888 empty,
889 empty,
890 ctx.scale,
891 use_qk_l2norm_in_kernel=False,
892 )
894 if send_work is not None:
895 send_work.wait()
896 return dq, dk, dv, dg, dbeta, None, None, None, None, None
899def _gdn_state_p2p_triton(
900 query: torch.Tensor,
901 key: torch.Tensor,
902 value: torch.Tensor,
903 g: torch.Tensor,
904 beta: torch.Tensor,
905 cp_mesh: DeviceMesh,
906 cp_rank: int,
907 cp_size: int,
908) -> torch.Tensor:
909 """Run fused local GDN with an affine state wavefront."""
910 if cp_size == 1:
911 output, _ = chunk_gated_delta_rule(
912 query,
913 key,
914 value,
915 g=g,
916 beta=beta,
917 output_final_state=False,
918 use_qk_l2norm_in_kernel=True,
919 backend="triton",
920 )
921 return output
923 query = _l2norm_torch(query)
924 key = _l2norm_torch(key)
925 prev_rank = _global_peer_rank(cp_mesh, cp_rank - 1) if cp_rank > 0 else -1
926 next_rank = (
927 _global_peer_rank(cp_mesh, cp_rank + 1) if cp_rank < cp_size - 1 else -1
928 )
929 return _GDNStateP2PTritonFunction.apply(
930 query,
931 key,
932 value,
933 g,
934 beta,
935 cp_rank,
936 cp_size,
937 cp_mesh.get_group(),
938 prev_rank,
939 next_rank,
940 )
943def _differentiable_all_to_all_shard(
944 tensor: torch.Tensor,
945 device_mesh: DeviceMesh,
946 *,
947 split_dim: int,
948 concat_dim: int,
949) -> torch.Tensor:
950 """Split local data on ``split_dim`` and concatenate peers on ``concat_dim``.
952 This is the local-tensor equivalent of DTensor ``Shard(concat_dim) ->
953 Shard(split_dim)`` redistribution for a 1-D mesh. It uses platform-level
954 differentiable all-to-all directly to avoid wrapping each activation in a
955 temporary DTensor.
956 """
957 split_count = device_mesh.size()
958 if split_count == 1:
959 return tensor
961 original_shape = tuple(tensor.shape)
962 dim_size = original_shape[split_dim]
963 if dim_size % split_count != 0:
964 raise ValueError(
965 f"linear attention all-to-all split dim {split_dim} with size "
966 f"{dim_size} must be divisible by cp_size {split_count}."
967 )
969 split_size = dim_size // split_count
970 final_shape = list(original_shape)
971 if split_dim != concat_dim:
972 final_shape[split_dim] = split_size
973 final_shape[concat_dim] = final_shape[concat_dim] * split_count
974 final_shape = tuple(final_shape)
976 reshape_dims = list(original_shape)
977 reshape_dims[split_dim] = split_count
978 reshape_dims.insert(split_dim + 1, split_size)
980 trans_dims = list(range(len(reshape_dims)))
981 trans_dims.remove(split_dim)
982 trans_dims.insert(0, split_dim)
984 a2a_input = tensor.reshape(reshape_dims).permute(trans_dims).contiguous()
985 reshape_shape = list(a2a_input.shape)
986 reshape_shape[0] = reshape_shape[0] * reshape_shape[1]
987 reshape_shape.pop(1)
988 a2a_input = a2a_input.reshape(reshape_shape)
990 a2a_input = a2a_input.contiguous()
991 split_len = a2a_input.shape[0] // split_count
992 input_splits = [split_len] * split_count
993 output_splits = [split_len] * split_count
994 output = platform.differentiable_all_to_all_single(
995 a2a_input,
996 input_splits,
997 output_splits,
998 group=device_mesh.get_group(),
999 )
1001 output_reshape = list(output.shape)
1002 output_reshape[0] = split_count
1003 output_reshape.insert(1, output.shape[0] // split_count)
1005 out_trans_dims = list(range(len(output_reshape)))
1006 first_dim = out_trans_dims.pop(0)
1007 if concat_dim >= len(out_trans_dims):
1008 out_trans_dims.append(first_dim)
1009 else:
1010 out_trans_dims.insert(concat_dim, first_dim)
1012 final_output = output.reshape(output_reshape).permute(out_trans_dims).contiguous()
1013 final_reshape = list(final_output.shape)
1014 if concat_dim < len(final_reshape) - 1:
1015 final_reshape[concat_dim] = (
1016 final_reshape[concat_dim] * final_reshape[concat_dim + 1]
1017 )
1018 final_reshape.pop(concat_dim + 1)
1020 return final_output.reshape(final_reshape).view(final_shape)
1023class LinearAttentionUlyssesCPWrapper(nn.Module):
1024 """Pure-Ulysses CP execution wrapper for a Qwen3.5 Gated DeltaNet module.
1026 Parameters stay owned by the original module. The wrapper only changes the
1027 execution layout:
1029 ``[B, S_local, full_heads] -> [B, S_full, local_heads] ->
1030 [B, S_local, full_heads]``.
1031 """
1033 def __init__(
1034 self,
1035 module: nn.Module,
1036 device_mesh: DeviceMesh,
1037 *,
1038 backend: str = "eager",
1039 ):
1040 super().__init__()
1041 self.module = module
1042 self.gdn_backend = backend
1043 self.cp_mesh = _ensure_1d(device_mesh)
1044 self.cp_size = self.cp_mesh.size()
1045 self.cp_rank = self.cp_mesh.get_local_rank()
1046 self.seq_dim = 1
1047 self.head_dim = 2
1048 self._validate_module()
1050 def _validate_module(self) -> None:
1051 if self.cp_size <= 1:
1052 return
1053 if self.module.num_k_heads % self.cp_size != 0:
1054 raise ValueError(
1055 f"linear attention num_k_heads ({self.module.num_k_heads}) must be "
1056 f"divisible by cp_size ({self.cp_size}) for Ulysses CP."
1057 )
1058 if self.module.num_v_heads % self.cp_size != 0:
1059 raise ValueError(
1060 f"linear attention num_v_heads ({self.module.num_v_heads}) must be "
1061 f"divisible by cp_size ({self.cp_size}) for Ulysses CP."
1062 )
1064 def _seq_to_head(self, tensor: torch.Tensor) -> torch.Tensor:
1065 return _differentiable_all_to_all_shard(
1066 tensor,
1067 self.cp_mesh,
1068 split_dim=self.head_dim,
1069 concat_dim=self.seq_dim,
1070 )
1072 def _head_to_seq(self, tensor: torch.Tensor) -> torch.Tensor:
1073 return _differentiable_all_to_all_shard(
1074 tensor,
1075 self.cp_mesh,
1076 split_dim=self.seq_dim,
1077 concat_dim=self.head_dim,
1078 )
1080 def _seq_to_head_qkvba(
1081 self,
1082 q_proj: torch.Tensor,
1083 k_proj: torch.Tensor,
1084 v_proj: torch.Tensor,
1085 b: torch.Tensor,
1086 a: torch.Tensor,
1087 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
1088 """Pack Q/K/V/B/A by CP rank and run a single seq-to-head all-to-all."""
1089 if self.cp_size == 1:
1090 return q_proj, k_proj, v_proj, b, a
1092 base = self.module
1093 local_key_dim = base.key_dim // self.cp_size
1094 local_value_dim = base.value_dim // self.cp_size
1095 local_num_v_heads = base.num_v_heads // self.cp_size
1097 q_chunks = torch.split(q_proj, local_key_dim, dim=-1)
1098 k_chunks = torch.split(k_proj, local_key_dim, dim=-1)
1099 v_chunks = torch.split(v_proj, local_value_dim, dim=-1)
1100 b_chunks = torch.split(b, local_num_v_heads, dim=-1)
1101 a_chunks = torch.split(a, local_num_v_heads, dim=-1)
1102 rank_major_chunks = [
1103 torch.cat(chunks, dim=-1)
1104 for chunks in zip(q_chunks, k_chunks, v_chunks, b_chunks, a_chunks)
1105 ]
1106 packed = torch.cat(rank_major_chunks, dim=-1).contiguous()
1107 packed = self._seq_to_head(packed)
1108 return torch.split(
1109 packed,
1110 [
1111 local_key_dim,
1112 local_key_dim,
1113 local_value_dim,
1114 local_num_v_heads,
1115 local_num_v_heads,
1116 ],
1117 dim=-1,
1118 )
1120 def _local_conv_weight(self) -> torch.Tensor:
1121 return _slice_qkv_local_cp(
1122 self.module.conv1d.weight,
1123 key_dim=self.module.key_dim,
1124 value_dim=self.module.value_dim,
1125 dim=0,
1126 cp_rank=self.cp_rank,
1127 cp_size=self.cp_size,
1128 )
1130 def _local_conv_bias(self) -> Optional[torch.Tensor]:
1131 bias = self.module.conv1d.bias
1132 if bias is None:
1133 return None
1134 return _slice_qkv_local_cp(
1135 bias,
1136 key_dim=self.module.key_dim,
1137 value_dim=self.module.value_dim,
1138 dim=0,
1139 cp_rank=self.cp_rank,
1140 cp_size=self.cp_size,
1141 )
1143 def forward(
1144 self,
1145 hidden_states: torch.Tensor,
1146 attention_mask: Optional[torch.Tensor] = None,
1147 **kwargs,
1148 ) -> torch.Tensor:
1149 """Run Gated DeltaNet with pure Ulysses context parallel."""
1150 del kwargs
1151 hidden_states = _local_tensor_at_cp_boundary(hidden_states)
1153 base = self.module
1154 if attention_mask is not None and attention_mask.ndim == 2:
1155 hidden_states = hidden_states * attention_mask[:, :, None].to(
1156 hidden_states.dtype
1157 )
1159 bsz, local_seq_len, _ = hidden_states.shape
1160 mixed_qkv = base.in_proj_qkv(hidden_states)
1161 z = base.in_proj_z(hidden_states).reshape(
1162 bsz,
1163 local_seq_len,
1164 base.num_v_heads,
1165 base.head_v_dim,
1166 )
1167 b = base.in_proj_b(hidden_states)
1168 a = base.in_proj_a(hidden_states)
1170 q_proj, k_proj, v_proj = torch.split(
1171 mixed_qkv,
1172 [base.key_dim, base.key_dim, base.value_dim],
1173 dim=-1,
1174 )
1175 q_proj, k_proj, v_proj, b, a = self._seq_to_head_qkvba(
1176 q_proj, k_proj, v_proj, b, a
1177 )
1179 full_seq_len = q_proj.shape[1]
1180 local_key_dim = base.key_dim // self.cp_size
1181 local_value_dim = base.value_dim // self.cp_size
1182 local_num_k_heads = base.num_k_heads // self.cp_size
1183 local_num_v_heads = base.num_v_heads // self.cp_size
1184 local_conv_dim = local_key_dim * 2 + local_value_dim
1186 mixed_qkv = torch.cat((q_proj, k_proj, v_proj), dim=-1).transpose(1, 2)
1187 conv_out = F.conv1d(
1188 input=mixed_qkv,
1189 weight=self._local_conv_weight(),
1190 bias=self._local_conv_bias(),
1191 stride=base.conv1d.stride,
1192 padding=base.conv1d.padding,
1193 dilation=base.conv1d.dilation,
1194 groups=local_conv_dim,
1195 )
1196 mixed_qkv = F.silu(conv_out[:, :, :full_seq_len]).transpose(1, 2)
1198 query, key, value = torch.split(
1199 mixed_qkv,
1200 [local_key_dim, local_key_dim, local_value_dim],
1201 dim=-1,
1202 )
1203 query = query.reshape(bsz, full_seq_len, local_num_k_heads, base.head_k_dim)
1204 key = key.reshape(bsz, full_seq_len, local_num_k_heads, base.head_k_dim)
1205 value = value.reshape(bsz, full_seq_len, local_num_v_heads, base.head_v_dim)
1207 a_log = _slice_local_cp(base.A_log, 0, self.cp_rank, self.cp_size)
1208 dt_bias = _slice_local_cp(base.dt_bias, 0, self.cp_rank, self.cp_size)
1209 beta = b.sigmoid()
1210 g = -a_log.float().exp() * F.softplus(a.float() + dt_bias)
1212 if base.kv_groups > 1:
1213 query = query.repeat_interleave(base.kv_groups, dim=2)
1214 key = key.repeat_interleave(base.kv_groups, dim=2)
1216 core_attn_out, _ = chunk_gated_delta_rule(
1217 query,
1218 key,
1219 value,
1220 g=g,
1221 beta=beta,
1222 initial_state=None,
1223 output_final_state=False,
1224 use_qk_l2norm_in_kernel=True,
1225 backend=self.gdn_backend,
1226 )
1228 core_attn_out = self._head_to_seq(core_attn_out)
1229 core_attn_out = core_attn_out.reshape(-1, base.head_v_dim)
1230 z_flat = z.reshape(-1, base.head_v_dim)
1231 core_attn_out = base.norm(core_attn_out, z_flat)
1232 core_attn_out = core_attn_out.reshape(bsz, local_seq_len, base.value_dim)
1233 if hasattr(base, "out_proj_input"):
1234 core_attn_out = base.out_proj_input(core_attn_out)
1235 return base.out_proj(core_attn_out)
1238class LinearAttentionP2PCPWrapper(nn.Module):
1239 """Sequence-sharded GDN CP with an affine-summary state wavefront."""
1241 def __init__(
1242 self,
1243 module: nn.Module,
1244 device_mesh: DeviceMesh,
1245 *,
1246 backend: str = "eager",
1247 ):
1248 super().__init__()
1249 self.module = module
1250 self.gdn_backend = backend
1251 self.cp_mesh = _ensure_1d(device_mesh)
1252 self.cp_size = self.cp_mesh.size()
1253 self.cp_rank = self.cp_mesh.get_local_rank()
1254 self._validate_module()
1256 def _validate_module(self) -> None:
1257 """Validate the Conv1d requirements of the P2P CP path."""
1258 if self.gdn_backend == "triton" and (
1259 self.module.head_k_dim != 128 or self.module.head_v_dim != 128
1260 ):
1261 raise NotImplementedError(
1262 "linear attention P2P Triton backend requires "
1263 "head_k_dim=head_v_dim=128."
1264 )
1265 conv = self.module.conv1d
1266 if conv.stride != (1,):
1267 raise ValueError(
1268 "linear attention P2P CP currently supports only conv1d stride=1."
1269 )
1270 if conv.groups != self.module.conv_dim:
1271 raise ValueError(
1272 "linear attention P2P CP expects depthwise conv1d groups=conv_dim."
1273 )
1274 if (
1275 conv.in_channels != self.module.conv_dim
1276 or conv.out_channels != self.module.conv_dim
1277 ):
1278 raise ValueError(
1279 "linear attention P2P CP expects conv1d channels to match conv_dim."
1280 )
1282 def _conv1d_with_halo(self, mixed_qkv: torch.Tensor) -> torch.Tensor:
1283 """Run local Conv1d after exchanging only the previous-rank halo."""
1284 return _causal_conv1d_with_cp_halo(
1285 mixed_qkv,
1286 self.module.conv1d,
1287 self.cp_mesh,
1288 self.cp_rank,
1289 self.cp_size,
1290 )
1292 def forward(
1293 self,
1294 hidden_states: torch.Tensor,
1295 attention_mask: Optional[torch.Tensor] = None,
1296 **kwargs,
1297 ) -> torch.Tensor:
1298 """Run Gated DeltaNet on local sequence shards with recurrent-state P2P."""
1299 del kwargs
1300 hidden_states = _local_tensor_at_cp_boundary(hidden_states)
1302 base = self.module
1303 if attention_mask is not None and attention_mask.ndim == 2:
1304 hidden_states = hidden_states * attention_mask[:, :, None].to(
1305 hidden_states.dtype
1306 )
1308 bsz, local_seq_len, _ = hidden_states.shape
1309 mixed_qkv = base.in_proj_qkv(hidden_states)
1310 z = base.in_proj_z(hidden_states).reshape(
1311 bsz,
1312 local_seq_len,
1313 base.num_v_heads,
1314 base.head_v_dim,
1315 )
1316 b = base.in_proj_b(hidden_states)
1317 a = base.in_proj_a(hidden_states)
1319 mixed_qkv = self._conv1d_with_halo(mixed_qkv)
1320 query, key, value = torch.split(
1321 mixed_qkv,
1322 [base.key_dim, base.key_dim, base.value_dim],
1323 dim=-1,
1324 )
1325 query = query.reshape(bsz, local_seq_len, base.num_k_heads, base.head_k_dim)
1326 key = key.reshape(bsz, local_seq_len, base.num_k_heads, base.head_k_dim)
1327 value = value.reshape(bsz, local_seq_len, base.num_v_heads, base.head_v_dim)
1329 beta = b.sigmoid()
1330 g = -base.A_log.float().exp() * F.softplus(a.float() + base.dt_bias)
1332 if base.kv_groups > 1:
1333 query = query.repeat_interleave(base.kv_groups, dim=2)
1334 key = key.repeat_interleave(base.kv_groups, dim=2)
1336 if self.gdn_backend == "triton":
1337 if local_seq_len % 64 != 0:
1338 raise NotImplementedError(
1339 "linear attention P2P Triton backend requires each CP "
1340 f"rank's local sequence length ({local_seq_len}) to be "
1341 "divisible by 64."
1342 )
1343 if not is_triton_gdn_available(query, key, value, g, beta):
1344 raise RuntimeError(
1345 "linear attention P2P Triton backend requires an NPU "
1346 "input satisfying the fixed GDN contract and a validated "
1347 "triton-ascend 3.2.x installation."
1348 )
1349 core_attn_out = _gdn_state_p2p_triton(
1350 query,
1351 key,
1352 value,
1353 g,
1354 beta,
1355 self.cp_mesh,
1356 self.cp_rank,
1357 self.cp_size,
1358 )
1359 else:
1360 core_attn_out = _gdn_state_p2p_summary(
1361 query,
1362 key,
1363 value,
1364 g,
1365 beta,
1366 self.cp_mesh,
1367 self.cp_rank,
1368 self.cp_size,
1369 use_qk_l2norm_in_kernel=True,
1370 )
1372 core_attn_out = core_attn_out.reshape(-1, base.head_v_dim)
1373 z_flat = z.reshape(-1, base.head_v_dim)
1374 core_attn_out = base.norm(core_attn_out, z_flat)
1375 core_attn_out = core_attn_out.reshape(bsz, local_seq_len, base.value_dim)
1376 if hasattr(base, "out_proj_input"):
1377 core_attn_out = base.out_proj_input(core_attn_out)
1378 return base.out_proj(core_attn_out)
1381class LinearAttentionAllGatherCPWrapper(nn.Module):
1382 """Sequence-sharded GDN CP using all-gathered recurrent-state summaries."""
1384 def __init__(self, module: nn.Module, device_mesh: DeviceMesh):
1385 super().__init__()
1386 self.module = module
1387 self.cp_mesh = _ensure_1d(device_mesh)
1388 self.cp_size = self.cp_mesh.size()
1389 self.cp_rank = self.cp_mesh.get_local_rank()
1390 self._validate_module()
1392 def _validate_module(self) -> None:
1393 """Validate the Conv1d requirements of the all-gather CP path."""
1394 conv = self.module.conv1d
1395 if conv.stride != (1,):
1396 raise ValueError(
1397 "linear attention all-gather CP currently supports only "
1398 "conv1d stride=1."
1399 )
1400 if conv.groups != self.module.conv_dim:
1401 raise ValueError(
1402 "linear attention all-gather CP expects depthwise conv1d "
1403 "groups=conv_dim."
1404 )
1405 if (
1406 conv.in_channels != self.module.conv_dim
1407 or conv.out_channels != self.module.conv_dim
1408 ):
1409 raise ValueError(
1410 "linear attention all-gather CP expects conv1d channels to "
1411 "match conv_dim."
1412 )
1414 def _conv1d_with_halo(self, mixed_qkv: torch.Tensor) -> torch.Tensor:
1415 """Run local Conv1d after exchanging only the previous-rank halo."""
1416 return _causal_conv1d_with_cp_halo(
1417 mixed_qkv,
1418 self.module.conv1d,
1419 self.cp_mesh,
1420 self.cp_rank,
1421 self.cp_size,
1422 )
1424 def forward(
1425 self,
1426 hidden_states: torch.Tensor,
1427 attention_mask: Optional[torch.Tensor] = None,
1428 **kwargs,
1429 ) -> torch.Tensor:
1430 """Run Gated DeltaNet on local sequence shards with all-gather state summaries."""
1431 del kwargs
1432 hidden_states = _local_tensor_at_cp_boundary(hidden_states)
1434 base = self.module
1435 if attention_mask is not None and attention_mask.ndim == 2:
1436 hidden_states = hidden_states * attention_mask[:, :, None].to(
1437 hidden_states.dtype
1438 )
1440 bsz, local_seq_len, _ = hidden_states.shape
1441 mixed_qkv = base.in_proj_qkv(hidden_states)
1442 z = base.in_proj_z(hidden_states).reshape(
1443 bsz,
1444 local_seq_len,
1445 base.num_v_heads,
1446 base.head_v_dim,
1447 )
1448 b = base.in_proj_b(hidden_states)
1449 a = base.in_proj_a(hidden_states)
1451 mixed_qkv = self._conv1d_with_halo(mixed_qkv)
1452 query, key, value = torch.split(
1453 mixed_qkv,
1454 [base.key_dim, base.key_dim, base.value_dim],
1455 dim=-1,
1456 )
1457 query = query.reshape(bsz, local_seq_len, base.num_k_heads, base.head_k_dim)
1458 key = key.reshape(bsz, local_seq_len, base.num_k_heads, base.head_k_dim)
1459 value = value.reshape(bsz, local_seq_len, base.num_v_heads, base.head_v_dim)
1461 beta = b.sigmoid()
1462 g = -base.A_log.float().exp() * F.softplus(a.float() + base.dt_bias)
1464 if base.kv_groups > 1:
1465 query = query.repeat_interleave(base.kv_groups, dim=2)
1466 key = key.repeat_interleave(base.kv_groups, dim=2)
1468 core_attn_out = _gdn_state_all_gather(
1469 query,
1470 key,
1471 value,
1472 g,
1473 beta,
1474 self.cp_mesh,
1475 self.cp_rank,
1476 self.cp_size,
1477 use_qk_l2norm_in_kernel=True,
1478 )
1480 core_attn_out = core_attn_out.reshape(-1, base.head_v_dim)
1481 z_flat = z.reshape(-1, base.head_v_dim)
1482 core_attn_out = base.norm(core_attn_out, z_flat)
1483 core_attn_out = core_attn_out.reshape(bsz, local_seq_len, base.value_dim)
1484 if hasattr(base, "out_proj_input"):
1485 core_attn_out = base.out_proj_input(core_attn_out)
1486 return base.out_proj(core_attn_out)
1489class LinearAttentionContextParallel(ParallelStyle):
1490 """Apply context parallel execution to a Gated DeltaNet module."""
1492 def __init__(self, *, mode: str = "ulysses", backend: str = "eager") -> None:
1493 if mode not in {"ulysses", "p2p", "all_gather"}:
1494 raise NotImplementedError(
1495 "LinearAttentionContextParallel currently supports mode='ulysses', "
1496 "mode='p2p', and mode='all_gather'."
1497 )
1498 if backend not in {"eager", "triton"}:
1499 raise ValueError(
1500 "LinearAttentionContextParallel backend must be 'eager' or "
1501 f"'triton', got {backend!r}."
1502 )
1503 if mode == "all_gather" and backend == "triton":
1504 raise NotImplementedError(
1505 "linear attention all-gather CP does not yet support the "
1506 "Triton backend."
1507 )
1508 self.mode = mode
1509 self.backend = backend
1511 def apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
1512 """Patch ``module.forward`` with a linear-attention CP executor."""
1513 if self.mode == "ulysses":
1514 executor = LinearAttentionUlyssesCPWrapper(
1515 module,
1516 device_mesh,
1517 backend=self.backend,
1518 )
1519 elif self.mode == "all_gather":
1520 executor = LinearAttentionAllGatherCPWrapper(module, device_mesh)
1521 else:
1522 executor = LinearAttentionP2PCPWrapper(
1523 module,
1524 device_mesh,
1525 backend=self.backend,
1526 )
1527 object.__setattr__(module, "_hp_linear_attention_cp_executor", executor)
1528 object.__setattr__(module, "_hp_linear_attention_original_forward", module.forward)
1530 def _forward(*args, **kwargs):
1531 return executor(*args, **kwargs)
1533 object.__setattr__(module, "forward", _forward)
1534 return module