Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / gated_delta_net.py: 14%
525 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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# pylint: disable=forbidden-backend-import,missing-public-type-hints
18# pylint: disable=missing-public-docstring,not-callable
19# PyTorch autograd.Function intentionally defines framework-specific signatures.
20# pylint: disable=abstract-method,arguments-differ
21from __future__ import annotations
23__all__ = ["GatedDeltaNetP2PCP", "GatedDeltaNetUlyssesCP"]
25from typing import NamedTuple, Optional
27import torch
28import torch.distributed as dist
29from torch import nn
30from torch.nn import functional as F
31from torch.utils.checkpoint import checkpoint
33from hyper_parallel.core.context_parallel.context_parallel import (
34 _ensure_1d,
35)
36from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
37from hyper_parallel.core.dtensor.dtensor import DTensor
38from hyper_parallel.components.modules.gated_delta_net import (
39 chunk_gated_delta_rule,
40 is_triton_gdn_available,
41 torch_chunk_gated_delta_rule,
42)
43from hyper_parallel.core.utils import communication
46_GDN_BACKENDS = frozenset({"eager", "triton"})
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 _kv_groups(module: nn.Module) -> int:
106 """Return the Q/K replication factor used by the official HF GDN layer."""
107 return int(module.num_v_heads // module.num_k_heads)
110def _apply_local_padding_mask(
111 hidden_states: torch.Tensor,
112 attention_mask: Optional[torch.Tensor],
113) -> torch.Tensor:
114 """Apply the supported local 2-D padding mask without silent fallback."""
115 if attention_mask is None:
116 return hidden_states
117 if attention_mask.ndim != 2:
118 raise NotImplementedError(
119 "Qwen3.5 GDN Context Parallel supports only a local 2-D padding mask."
120 )
121 if tuple(attention_mask.shape) != tuple(hidden_states.shape[:2]):
122 raise ValueError(
123 "Qwen3.5 GDN Context Parallel expects attention_mask shape "
124 f"{tuple(hidden_states.shape[:2])}, got {tuple(attention_mask.shape)}."
125 )
126 return hidden_states * attention_mask[:, :, None].to(hidden_states.dtype)
129def _validate_gdn_activation(module: nn.Module) -> None:
130 """Reject model variants whose Conv activation is not implemented here."""
131 activation = getattr(module, "activation", "silu")
132 if activation not in ("silu", "swish"):
133 raise NotImplementedError(
134 "Qwen3.5 GDN Context Parallel currently supports only SiLU/Swish "
135 f"Conv activation, got {activation!r}."
136 )
139def _all_to_all_previous_rank_halo(
140 tail: torch.Tensor,
141 cp_mesh: DeviceMesh,
142 cp_rank: int,
143 cp_size: int,
144) -> torch.Tensor:
145 """Send a convolution halo only to the next rank using differentiable A2AV."""
146 if cp_size == 1:
147 return torch.zeros_like(tail)
149 cp_group = cp_mesh.get_group()
150 group_ranks = tuple(int(rank) for rank in dist.get_process_group_ranks(cp_group))
151 rank_list = tuple(int(rank) for rank in cp_mesh.rank_list)
152 rank_to_group_index = {rank: index for index, rank in enumerate(group_ranks)}
153 halo_width = tail.shape[1]
155 input_splits = [0] * cp_size
156 exchange_input = tail.permute(1, 0, 2).contiguous()
157 if cp_rank < cp_size - 1:
158 input_splits[rank_to_group_index[rank_list[cp_rank + 1]]] = halo_width
159 else:
160 exchange_input = exchange_input[:0]
162 output_splits = [0] * cp_size
163 if cp_rank > 0:
164 output_splits[rank_to_group_index[rank_list[cp_rank - 1]]] = halo_width
166 exchange_output = communication.differentiable_all_to_all_single(
167 exchange_input,
168 input_splits,
169 output_splits,
170 group=cp_group,
171 )
172 if cp_rank == 0:
173 return torch.zeros_like(tail) + exchange_output.sum().to(tail.dtype) * 0
174 return exchange_output.permute(1, 0, 2).contiguous()
177def _causal_conv1d_with_cp_halo(
178 mixed_qkv: torch.Tensor,
179 conv1d: nn.Conv1d,
180 cp_mesh: DeviceMesh,
181 cp_rank: int,
182 cp_size: int,
183) -> torch.Tensor:
184 """Run causal depthwise Conv1d with only the previous rank's boundary."""
185 kernel_size = conv1d.kernel_size[0]
186 dilation = conv1d.dilation[0]
187 halo_width = (kernel_size - 1) * dilation
188 if halo_width == 0 or cp_size == 1:
189 conv_out = conv1d(mixed_qkv.transpose(1, 2))
190 return F.silu(conv_out[:, :, : mixed_qkv.shape[1]]).transpose(1, 2)
192 if mixed_qkv.shape[1] < halo_width:
193 raise ValueError(
194 "linear attention CP conv halo requires local_seq_len >= "
195 f"{halo_width}, got {mixed_qkv.shape[1]}."
196 )
198 halo = _all_to_all_previous_rank_halo(
199 mixed_qkv[:, -halo_width:, :].contiguous(),
200 cp_mesh,
201 cp_rank,
202 cp_size,
203 )
204 conv_input = torch.cat((halo, mixed_qkv), dim=1).transpose(1, 2)
205 conv_out = F.conv1d(
206 input=conv_input,
207 weight=conv1d.weight,
208 bias=conv1d.bias,
209 stride=conv1d.stride,
210 padding=0,
211 dilation=conv1d.dilation,
212 groups=conv1d.groups,
213 )
214 return F.silu(conv_out).transpose(1, 2)
217def _l2norm_torch(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor:
218 """Match the pure torch GDN reference l2norm helper."""
219 return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
222class _GDNPreparedChunks(NamedTuple):
223 """Reusable chunk intermediates shared by state-summary CP modes."""
225 initial_dtype: torch.dtype
226 query: torch.Tensor
227 key: torch.Tensor
228 chunk_value: torch.Tensor
229 g: torch.Tensor
230 decay_mask: torch.Tensor
231 k_cumdecay: torch.Tensor
232 sequence_length: int
233 total_sequence_length: int
234 chunk_size: int
237def _prepare_gdn_chunks_for_summary(
238 query: torch.Tensor,
239 key: torch.Tensor,
240 value: torch.Tensor,
241 g: torch.Tensor,
242 beta: torch.Tensor,
243 *,
244 chunk_size: int = 64,
245 use_qk_l2norm_in_kernel: bool = False,
246) -> _GDNPreparedChunks:
247 """Prepare GDN chunk intermediates shared by summary and local output."""
248 initial_dtype = query.dtype
249 if use_qk_l2norm_in_kernel:
250 query = _l2norm_torch(query, dim=-1, eps=1e-6)
251 key = _l2norm_torch(key, dim=-1, eps=1e-6)
253 query, key, value, beta, g = [
254 x.transpose(1, 2).contiguous().to(torch.float32)
255 for x in (query, key, value, beta, g)
256 ]
258 sequence_length = key.shape[2]
259 pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
260 query = F.pad(query, (0, 0, 0, pad_size))
261 key = F.pad(key, (0, 0, 0, pad_size))
262 value = F.pad(value, (0, 0, 0, pad_size))
263 beta = F.pad(beta, (0, pad_size))
264 g = F.pad(g, (0, pad_size))
265 total_sequence_length = sequence_length + pad_size
267 query = query * (1 / (query.shape[-1] ** 0.5))
268 v_beta = value * beta.unsqueeze(-1)
269 k_beta = key * beta.unsqueeze(-1)
270 query, key, k_beta, v_beta = [
271 x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
272 for x in (query, key, k_beta, v_beta)
273 ]
274 g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
276 mask = torch.triu(
277 torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
278 diagonal=0,
279 )
280 g = g.cumsum(dim=-1)
281 decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
282 attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
283 for row_idx in range(1, chunk_size):
284 row = attn[..., row_idx, :row_idx].clone()
285 sub = attn[..., :row_idx, :row_idx].clone()
286 attn[..., row_idx, :row_idx] = row + (row.unsqueeze(-1) * sub).sum(-2)
287 attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
289 chunk_value = attn @ v_beta
290 k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
291 return _GDNPreparedChunks(
292 initial_dtype=initial_dtype,
293 query=query,
294 key=key,
295 chunk_value=chunk_value,
296 g=g,
297 decay_mask=decay_mask,
298 k_cumdecay=k_cumdecay,
299 sequence_length=sequence_length,
300 total_sequence_length=total_sequence_length,
301 chunk_size=chunk_size,
302 )
305def _compute_gdn_state_summary_from_prepared(
306 prepared: _GDNPreparedChunks,
307) -> tuple[torch.Tensor, torch.Tensor]:
308 """Compute ``state_out = M @ state_in + S`` from prepared GDN chunks."""
309 key = prepared.key
310 batch_size, num_heads, _, _, k_head_dim = key.shape
311 v_head_dim = prepared.chunk_value.shape[-1]
312 eye = torch.eye(k_head_dim, device=key.device, dtype=torch.float32).reshape(
313 1, 1, k_head_dim, k_head_dim
314 )
315 state_ext = torch.zeros(
316 batch_size,
317 num_heads,
318 k_head_dim,
319 v_head_dim,
320 device=key.device,
321 dtype=torch.float32,
322 )
323 transition = eye.expand(batch_size, num_heads, -1, -1).clone()
325 for chunk_idx in range(key.shape[2]):
326 key_i = key[:, :, chunk_idx]
327 value_i = prepared.chunk_value[:, :, chunk_idx]
328 w_i = prepared.k_cumdecay[:, :, chunk_idx]
329 g_i = prepared.g[:, :, chunk_idx]
330 decay = g_i[:, :, -1].exp()
331 key_decay = key_i * (g_i[:, :, -1, None] - g_i).exp()[..., None]
333 transition_i = (
334 decay[:, :, None, None] * eye
335 - key_decay.transpose(-1, -2) @ w_i
336 )
337 state_ext_i = key_decay.transpose(-1, -2) @ value_i
338 state_ext = transition_i @ state_ext + state_ext_i
339 transition = transition_i @ transition
341 return state_ext, transition
344def _checkpoint_gdn_state_summary(
345 prepared: _GDNPreparedChunks,
346) -> tuple[torch.Tensor, torch.Tensor]:
347 """Compute a state summary without retaining its per-chunk autograd graph."""
348 if not torch.is_grad_enabled():
349 return _compute_gdn_state_summary_from_prepared(prepared)
351 def recompute(
352 key: torch.Tensor,
353 chunk_value: torch.Tensor,
354 g: torch.Tensor,
355 k_cumdecay: torch.Tensor,
356 ) -> tuple[torch.Tensor, torch.Tensor]:
357 """Rebuild a prepared view from explicit checkpoint inputs."""
358 checkpoint_prepared = prepared._replace(
359 key=key,
360 chunk_value=chunk_value,
361 g=g,
362 k_cumdecay=k_cumdecay,
363 )
364 return _compute_gdn_state_summary_from_prepared(checkpoint_prepared)
366 return checkpoint(
367 recompute,
368 prepared.key,
369 prepared.chunk_value,
370 prepared.g,
371 prepared.k_cumdecay,
372 use_reentrant=False,
373 preserve_rng_state=False,
374 )
377def _run_prepared_gdn_chunks(
378 prepared: _GDNPreparedChunks,
379 initial_state: Optional[torch.Tensor],
380) -> torch.Tensor:
381 """Run local GDN output using already prepared chunk intermediates."""
382 query = prepared.query
383 key = prepared.key
384 chunk_value = prepared.chunk_value
385 batch_size, num_heads, _, _, k_head_dim = key.shape
386 v_head_dim = chunk_value.shape[-1]
387 recurrent_state = (
388 torch.zeros(
389 batch_size,
390 num_heads,
391 k_head_dim,
392 v_head_dim,
393 device=chunk_value.device,
394 dtype=chunk_value.dtype,
395 )
396 if initial_state is None
397 else initial_state.to(chunk_value)
398 )
399 core_attn_out = torch.zeros_like(chunk_value)
401 for chunk_idx in range(0, prepared.total_sequence_length // prepared.chunk_size):
402 q_i = query[:, :, chunk_idx]
403 k_i = key[:, :, chunk_idx]
404 v_i = chunk_value[:, :, chunk_idx]
405 attn = q_i @ k_i.transpose(-1, -2) * prepared.decay_mask[:, :, chunk_idx]
406 v_prime = prepared.k_cumdecay[:, :, chunk_idx] @ recurrent_state
407 v_new = v_i - v_prime
408 attn_inter = (
409 q_i * prepared.g[:, :, chunk_idx, :, None].exp()
410 ) @ recurrent_state
411 core_attn_out[:, :, chunk_idx] = attn_inter + attn @ v_new
412 recurrent_state = (
413 recurrent_state * prepared.g[:, :, chunk_idx, -1, None, None].exp()
414 + (
415 k_i
416 * (
417 prepared.g[:, :, chunk_idx, -1, None]
418 - prepared.g[:, :, chunk_idx]
419 ).exp()[..., None]
420 ).transpose(-1, -2) @ v_new
421 )
423 core_attn_out = core_attn_out.reshape(
424 core_attn_out.shape[0],
425 core_attn_out.shape[1],
426 -1,
427 core_attn_out.shape[-1],
428 )
429 core_attn_out = core_attn_out[:, :, :prepared.sequence_length]
430 return core_attn_out.transpose(1, 2).contiguous().to(prepared.initial_dtype)
433class _RecvInitialStateP2PFunction(torch.autograd.Function):
434 """Receive the recurrent initial state; send its gradient in backward."""
436 @staticmethod
437 def forward( # pylint: disable=arguments-differ
438 ctx,
439 anchor: torch.Tensor,
440 cp_group,
441 prev_rank: int,
442 state_shape: tuple[int, ...],
443 ) -> torch.Tensor:
444 """Receive the initial state from the preceding CP rank."""
445 state = torch.empty(state_shape, device=anchor.device, dtype=torch.float32)
446 dist.recv(state, src=prev_rank, group=cp_group)
447 ctx.cp_group = cp_group
448 ctx.prev_rank = prev_rank
449 return state
451 @staticmethod
452 def backward(ctx, grad_state: Optional[torch.Tensor]):
453 """Return the incoming state gradient to the preceding CP rank."""
454 if grad_state is None:
455 raise RuntimeError("linear attention P2P backward missing initial-state grad.")
456 dist.send(grad_state.contiguous(), dst=ctx.prev_rank, group=ctx.cp_group)
457 return None, None, None, None
460class _SendFinalStateP2PFunction(torch.autograd.Function):
461 """Send the recurrent final state; receive its gradient in backward."""
463 @staticmethod
464 def forward( # pylint: disable=arguments-differ
465 ctx,
466 final_state: torch.Tensor,
467 cp_group,
468 next_rank: int,
469 ) -> torch.Tensor:
470 """Send the final state to the succeeding CP rank."""
471 dist.send(final_state.contiguous(), dst=next_rank, group=cp_group)
472 ctx.cp_group = cp_group
473 ctx.next_rank = next_rank
474 ctx.state_shape = tuple(final_state.shape)
475 ctx.state_dtype = final_state.dtype
476 return final_state.new_zeros(())
478 @staticmethod
479 def backward(ctx, grad_token: torch.Tensor):
480 """Receive the final-state gradient from the succeeding CP rank."""
481 grad_state = torch.empty(
482 ctx.state_shape,
483 device=grad_token.device,
484 dtype=ctx.state_dtype,
485 )
486 dist.recv(grad_state, src=ctx.next_rank, group=ctx.cp_group)
487 return grad_state, None, None
490def _apply_gdn_state_summary(
491 state_ext: torch.Tensor,
492 transition: torch.Tensor,
493 initial_state: Optional[torch.Tensor],
494) -> torch.Tensor:
495 """Apply ``state_out = M @ state_in + S`` to an incoming GDN state."""
496 if initial_state is None:
497 return state_ext
498 return transition @ initial_state.to(transition) + state_ext
501def _gdn_state_p2p_summary(
502 query: torch.Tensor,
503 key: torch.Tensor,
504 value: torch.Tensor,
505 g: torch.Tensor,
506 beta: torch.Tensor,
507 cp_mesh: DeviceMesh,
508 cp_rank: int,
509 cp_size: int,
510 *,
511 chunk_size: int,
512 use_qk_l2norm_in_kernel: bool,
513) -> torch.Tensor:
514 """Run local GDN with an affine-summary state wavefront.
516 Every rank prepares its local chunks and state transition in parallel.
517 The rank-ordered critical path then contains only ``M @ state + S`` and
518 the small state transfer. Token outputs retain the ordinary PyTorch graph,
519 while the two custom autograd boundaries reverse the state communication.
520 """
521 if cp_size == 1:
522 core_attn_out, _ = torch_chunk_gated_delta_rule(
523 query,
524 key,
525 value,
526 g=g,
527 beta=beta,
528 initial_state=None,
529 output_final_state=False,
530 chunk_size=chunk_size,
531 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
532 )
533 return core_attn_out
535 cp_group = cp_mesh.get_group()
536 prev_rank = _global_peer_rank(cp_mesh, cp_rank - 1) if cp_rank > 0 else -1
537 next_rank = _global_peer_rank(cp_mesh, cp_rank + 1) if cp_rank < cp_size - 1 else -1
538 state_shape = (query.shape[0], value.shape[2], query.shape[3], value.shape[3])
540 prepared = _prepare_gdn_chunks_for_summary(
541 query,
542 key,
543 value,
544 g,
545 beta,
546 chunk_size=chunk_size,
547 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
548 )
549 state_ext = None
550 transition = None
551 if cp_rank < cp_size - 1:
552 state_ext, transition = _checkpoint_gdn_state_summary(prepared)
554 initial_state = None
555 if cp_rank > 0:
556 initial_state = _RecvInitialStateP2PFunction.apply(
557 query,
558 cp_group,
559 prev_rank,
560 state_shape,
561 )
563 send_token = None
564 if cp_rank < cp_size - 1:
565 final_state = _apply_gdn_state_summary(
566 state_ext,
567 transition,
568 initial_state,
569 )
570 send_token = _SendFinalStateP2PFunction.apply(final_state, cp_group, next_rank)
572 core_attn_out = _run_prepared_gdn_chunks(prepared, initial_state)
573 if send_token is not None:
574 core_attn_out = core_attn_out + send_token.to(core_attn_out.dtype) * 0
576 return core_attn_out
579class _GDNStateP2PTritonFunction(torch.autograd.Function):
580 """Pipeline fused affine GDN states over sequence-sharded CP ranks."""
582 @staticmethod
583 def forward( # pylint: disable=arguments-differ,too-many-locals
584 ctx,
585 query: torch.Tensor,
586 key: torch.Tensor,
587 value: torch.Tensor,
588 g: torch.Tensor,
589 beta: torch.Tensor,
590 cp_rank: int,
591 cp_size: int,
592 cp_group,
593 prev_rank: int,
594 next_rank: int,
595 ) -> torch.Tensor:
596 """Run fused local GDN and forward its affine state across CP ranks."""
597 from hyper_parallel.components.functional.gated_delta_net import ( # pylint: disable=import-outside-toplevel
598 chunk_gated_delta_rule_fwd_apply_state_saved,
599 chunk_gated_delta_rule_fwd_output_saved,
600 chunk_gated_delta_rule_fwd_prepare_saved,
601 )
602 from hyper_parallel.components.functional.gated_delta_net_state_summary import ( # pylint: disable=import-outside-toplevel
603 apply_gdn_state_summary,
604 chunk_gated_delta_rule_state_summary_fwd,
605 )
607 (
608 query_norm,
609 key_norm,
610 _,
611 _,
612 g_cumsum,
613 matrix_a,
614 w,
615 u,
616 scale,
617 ) = chunk_gated_delta_rule_fwd_prepare_saved(
618 query,
619 key,
620 value,
621 g,
622 beta,
623 use_qk_l2norm_in_kernel=False,
624 )
625 initial_state = None
626 recv_buffer = None
627 recv_work = None
628 if cp_rank > 0:
629 recv_buffer = torch.empty(
630 (query.shape[0], query.shape[2], query.shape[3], value.shape[3]),
631 device=query.device,
632 dtype=torch.float32,
633 )
634 recv_work = dist.irecv(recv_buffer, src=prev_rank, group=cp_group)
636 state_ext = None
637 transition = None
638 if cp_rank < cp_size - 1:
639 state_ext, transition = chunk_gated_delta_rule_state_summary_fwd(
640 key_norm,
641 w,
642 u,
643 g_cumsum,
644 )
646 if recv_work is not None:
647 recv_work.wait()
648 initial_state = recv_buffer
650 send_buffer = None
651 send_work = None
652 if cp_rank < cp_size - 1:
653 send_buffer = apply_gdn_state_summary(
654 state_ext,
655 transition,
656 initial_state,
657 ).contiguous()
658 send_work = dist.isend(send_buffer, dst=next_rank, group=cp_group)
660 h, v_new, _ = chunk_gated_delta_rule_fwd_apply_state_saved(
661 key_norm,
662 g_cumsum,
663 w,
664 u,
665 initial_state=initial_state,
666 output_final_state=False,
667 )
668 output = chunk_gated_delta_rule_fwd_output_saved(
669 query_norm,
670 key_norm,
671 g_cumsum,
672 h,
673 v_new,
674 scale,
675 ).to(query.dtype)
677 if send_work is not None:
678 send_work.wait()
680 empty = query.new_empty(0)
681 ctx.save_for_backward(
682 query_norm,
683 key_norm,
684 value,
685 g_cumsum,
686 beta,
687 matrix_a,
688 initial_state if initial_state is not None else empty,
689 transition if transition is not None else empty,
690 )
691 ctx.has_initial_state = initial_state is not None
692 ctx.cp_rank = cp_rank
693 ctx.cp_size = cp_size
694 ctx.cp_group = cp_group
695 ctx.prev_rank = prev_rank
696 ctx.next_rank = next_rank
697 ctx.scale = scale
698 return output
700 @staticmethod
701 def backward(ctx, grad_output: torch.Tensor): # pylint: disable=too-many-locals
702 """Backpropagate local GDN tensors and the state gradient wavefront."""
703 from hyper_parallel.components.functional.gated_delta_net import ( # pylint: disable=import-outside-toplevel
704 chunk_gated_delta_rule_bwd_finish_saved,
705 chunk_gated_delta_rule_bwd_prepare_saved,
706 chunk_gated_delta_rule_bwd_state_saved,
707 )
708 from hyper_parallel.components.functional.gated_delta_net_state_summary import ( # pylint: disable=import-outside-toplevel
709 apply_gdn_state_gradient_summary,
710 chunk_gated_delta_rule_state_gradient_summary_bwd,
711 )
713 (
714 query,
715 key,
716 value,
717 g_cumsum,
718 beta,
719 matrix_a,
720 initial_state,
721 transition,
722 ) = ctx.saved_tensors
723 if not ctx.has_initial_state:
724 initial_state = None
726 w, h, v_new, dv = chunk_gated_delta_rule_bwd_prepare_saved(
727 query,
728 key,
729 value,
730 g_cumsum,
731 beta,
732 matrix_a,
733 initial_state,
734 grad_output,
735 ctx.scale,
736 )
738 grad_state_ext = None
739 if ctx.cp_rank > 0:
740 grad_state_ext = chunk_gated_delta_rule_state_gradient_summary_bwd(
741 query,
742 key,
743 w,
744 g_cumsum,
745 grad_output,
746 dv,
747 ctx.scale,
748 )
750 grad_final_state = None
751 recv_work = None
752 if ctx.cp_rank < ctx.cp_size - 1:
753 recv_buffer = torch.empty(
754 (query.shape[0], query.shape[2], query.shape[3], value.shape[3]),
755 device=grad_output.device,
756 dtype=torch.float32,
757 )
758 recv_work = dist.irecv(
759 recv_buffer,
760 src=ctx.next_rank,
761 group=ctx.cp_group,
762 )
763 if recv_work is not None:
764 recv_work.wait()
765 grad_final_state = recv_buffer
767 send_buffer = None
768 send_work = None
769 if ctx.cp_rank > 0:
770 send_buffer = apply_gdn_state_gradient_summary(
771 grad_state_ext,
772 transition,
773 grad_final_state,
774 ).contiguous()
775 send_work = dist.isend(
776 send_buffer,
777 dst=ctx.prev_rank,
778 group=ctx.cp_group,
779 )
781 dh, _, dv = chunk_gated_delta_rule_bwd_state_saved(
782 query,
783 key,
784 g_cumsum,
785 w,
786 initial_state,
787 grad_final_state,
788 grad_output,
789 dv,
790 ctx.scale,
791 )
792 empty = query.new_empty(0)
793 dq, dk, dv, dg, dbeta = chunk_gated_delta_rule_bwd_finish_saved(
794 query,
795 key,
796 query,
797 key,
798 value,
799 g_cumsum,
800 beta,
801 matrix_a,
802 w,
803 h,
804 v_new,
805 dv,
806 grad_output,
807 dh,
808 empty,
809 empty,
810 ctx.scale,
811 use_qk_l2norm_in_kernel=False,
812 )
814 if send_work is not None:
815 send_work.wait()
816 return dq, dk, dv, dg, dbeta, None, None, None, None, None
819def _gdn_state_p2p_triton(
820 query: torch.Tensor,
821 key: torch.Tensor,
822 value: torch.Tensor,
823 g: torch.Tensor,
824 beta: torch.Tensor,
825 cp_mesh: DeviceMesh,
826 cp_rank: int,
827 cp_size: int,
828) -> torch.Tensor:
829 """Run fused local GDN with an affine state wavefront."""
830 if cp_size == 1:
831 output, _ = chunk_gated_delta_rule(
832 query,
833 key,
834 value,
835 g=g,
836 beta=beta,
837 output_final_state=False,
838 use_qk_l2norm_in_kernel=True,
839 backend="triton",
840 )
841 return output
843 query = _l2norm_torch(query)
844 key = _l2norm_torch(key)
845 prev_rank = _global_peer_rank(cp_mesh, cp_rank - 1) if cp_rank > 0 else -1
846 next_rank = (
847 _global_peer_rank(cp_mesh, cp_rank + 1) if cp_rank < cp_size - 1 else -1
848 )
849 return _GDNStateP2PTritonFunction.apply(
850 query,
851 key,
852 value,
853 g,
854 beta,
855 cp_rank,
856 cp_size,
857 cp_mesh.get_group(),
858 prev_rank,
859 next_rank,
860 )
863def _differentiable_all_to_all_shard(
864 tensor: torch.Tensor,
865 device_mesh: DeviceMesh,
866 *,
867 split_dim: int,
868 concat_dim: int,
869) -> torch.Tensor:
870 """Split local data on ``split_dim`` and concatenate peers on ``concat_dim``.
872 This is the local-tensor equivalent of DTensor ``Shard(concat_dim) ->
873 Shard(split_dim)`` redistribution for a 1-D mesh. It uses a
874 differentiable all-to-all directly to avoid wrapping each activation in a
875 temporary DTensor.
876 """
877 split_count = device_mesh.size()
878 if split_count == 1:
879 return tensor
881 original_shape = tuple(tensor.shape)
882 dim_size = original_shape[split_dim]
883 if dim_size % split_count != 0:
884 raise ValueError(
885 f"linear attention all-to-all split dim {split_dim} with size "
886 f"{dim_size} must be divisible by cp_size {split_count}."
887 )
889 split_size = dim_size // split_count
890 final_shape = list(original_shape)
891 if split_dim != concat_dim:
892 final_shape[split_dim] = split_size
893 final_shape[concat_dim] = final_shape[concat_dim] * split_count
894 final_shape = tuple(final_shape)
896 reshape_dims = list(original_shape)
897 reshape_dims[split_dim] = split_count
898 reshape_dims.insert(split_dim + 1, split_size)
900 trans_dims = list(range(len(reshape_dims)))
901 trans_dims.remove(split_dim)
902 trans_dims.insert(0, split_dim)
904 a2a_input = tensor.reshape(reshape_dims).permute(trans_dims).contiguous()
905 reshape_shape = list(a2a_input.shape)
906 reshape_shape[0] = reshape_shape[0] * reshape_shape[1]
907 reshape_shape.pop(1)
908 a2a_input = a2a_input.reshape(reshape_shape)
910 a2a_input = a2a_input.contiguous()
911 split_len = a2a_input.shape[0] // split_count
912 input_splits = [split_len] * split_count
913 output_splits = [split_len] * split_count
914 output = communication.differentiable_all_to_all_single(
915 a2a_input,
916 input_splits,
917 output_splits,
918 group=device_mesh.get_group(),
919 )
921 output_reshape = list(output.shape)
922 output_reshape[0] = split_count
923 output_reshape.insert(1, output.shape[0] // split_count)
925 out_trans_dims = list(range(len(output_reshape)))
926 first_dim = out_trans_dims.pop(0)
927 if concat_dim >= len(out_trans_dims):
928 out_trans_dims.append(first_dim)
929 else:
930 out_trans_dims.insert(concat_dim, first_dim)
932 final_output = output.reshape(output_reshape).permute(out_trans_dims).contiguous()
933 final_reshape = list(final_output.shape)
934 if concat_dim < len(final_reshape) - 1:
935 final_reshape[concat_dim] = (
936 final_reshape[concat_dim] * final_reshape[concat_dim + 1]
937 )
938 final_reshape.pop(concat_dim + 1)
940 return final_output.reshape(final_reshape).view(final_shape)
943class GatedDeltaNetUlyssesCP(nn.Module):
944 """Pure-Ulysses CP execution wrapper for a Qwen3.5 Gated DeltaNet module.
946 Parameters stay owned by the original module. The wrapper only changes the
947 execution layout:
949 ``[B, S_local, full_heads] -> [B, S_full, local_heads] ->
950 [B, S_local, full_heads]``.
951 """
953 def __init__(
954 self,
955 module: nn.Module,
956 device_mesh: DeviceMesh,
957 *,
958 backend: str = "eager",
959 chunk_size: int = 64,
960 ):
961 super().__init__()
962 self.module = module
963 backend = backend.lower()
964 if backend not in _GDN_BACKENDS:
965 raise ValueError(
966 f"unsupported GDN backend {backend!r}; "
967 f"expected one of {sorted(_GDN_BACKENDS)}."
968 )
969 if chunk_size <= 0:
970 raise ValueError(f"chunk_size must be positive, got {chunk_size}.")
971 if backend == "triton" and chunk_size != 64:
972 raise NotImplementedError("GDN Triton backend requires chunk_size=64.")
973 self.gdn_backend = backend
974 self.chunk_size = chunk_size
975 self.cp_mesh = _ensure_1d(device_mesh)
976 self.cp_size = self.cp_mesh.size()
977 self.cp_rank = self.cp_mesh.get_local_rank()
978 self.seq_dim = 1
979 self.head_dim = 2
980 self._validate_module()
982 def _validate_module(self) -> None:
983 """Validate the wrapped GDN module and Ulysses layout constraints."""
984 required = (
985 "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a",
986 "conv1d", "A_log", "dt_bias", "norm", "out_proj",
987 "num_k_heads", "num_v_heads", "head_k_dim", "head_v_dim",
988 "key_dim", "value_dim", "conv_dim",
989 )
990 missing = [name for name in required if not hasattr(self.module, name)]
991 if missing:
992 raise TypeError(
993 "Qwen3.5 GDN layer is missing required attributes: "
994 + ", ".join(missing)
995 )
996 if self.cp_size <= 1:
997 return
998 if self.module.num_k_heads % self.cp_size != 0:
999 raise ValueError(
1000 f"linear attention num_k_heads ({self.module.num_k_heads}) must be "
1001 f"divisible by cp_size ({self.cp_size}) for Ulysses CP."
1002 )
1003 if self.module.num_v_heads % self.cp_size != 0:
1004 raise ValueError(
1005 f"linear attention num_v_heads ({self.module.num_v_heads}) must be "
1006 f"divisible by cp_size ({self.cp_size}) for Ulysses CP."
1007 )
1008 if self.module.num_v_heads % self.module.num_k_heads:
1009 raise ValueError("GDN num_v_heads must be divisible by num_k_heads.")
1010 _validate_gdn_activation(self.module)
1012 def _seq_to_head(self, tensor: torch.Tensor) -> torch.Tensor:
1013 return _differentiable_all_to_all_shard(
1014 tensor,
1015 self.cp_mesh,
1016 split_dim=self.head_dim,
1017 concat_dim=self.seq_dim,
1018 )
1020 def _head_to_seq(self, tensor: torch.Tensor) -> torch.Tensor:
1021 return _differentiable_all_to_all_shard(
1022 tensor,
1023 self.cp_mesh,
1024 split_dim=self.seq_dim,
1025 concat_dim=self.head_dim,
1026 )
1028 def _seq_to_head_qkvba(
1029 self,
1030 q_proj: torch.Tensor,
1031 k_proj: torch.Tensor,
1032 v_proj: torch.Tensor,
1033 b: torch.Tensor,
1034 a: torch.Tensor,
1035 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
1036 """Pack Q/K/V/B/A by CP rank and run a single seq-to-head all-to-all."""
1037 if self.cp_size == 1:
1038 return q_proj, k_proj, v_proj, b, a
1040 base = self.module
1041 local_key_dim = base.key_dim // self.cp_size
1042 local_value_dim = base.value_dim // self.cp_size
1043 local_num_v_heads = base.num_v_heads // self.cp_size
1045 q_chunks = torch.split(q_proj, local_key_dim, dim=-1)
1046 k_chunks = torch.split(k_proj, local_key_dim, dim=-1)
1047 v_chunks = torch.split(v_proj, local_value_dim, dim=-1)
1048 b_chunks = torch.split(b, local_num_v_heads, dim=-1)
1049 a_chunks = torch.split(a, local_num_v_heads, dim=-1)
1050 rank_major_chunks = [
1051 torch.cat(chunks, dim=-1)
1052 for chunks in zip(q_chunks, k_chunks, v_chunks, b_chunks, a_chunks)
1053 ]
1054 packed = torch.cat(rank_major_chunks, dim=-1).contiguous()
1055 packed = self._seq_to_head(packed)
1056 return torch.split(
1057 packed,
1058 [
1059 local_key_dim,
1060 local_key_dim,
1061 local_value_dim,
1062 local_num_v_heads,
1063 local_num_v_heads,
1064 ],
1065 dim=-1,
1066 )
1068 def _local_conv_weight(self) -> torch.Tensor:
1069 return _slice_qkv_local_cp(
1070 self.module.conv1d.weight,
1071 key_dim=self.module.key_dim,
1072 value_dim=self.module.value_dim,
1073 dim=0,
1074 cp_rank=self.cp_rank,
1075 cp_size=self.cp_size,
1076 )
1078 def _local_conv_bias(self) -> Optional[torch.Tensor]:
1079 bias = self.module.conv1d.bias
1080 if bias is None:
1081 return None
1082 return _slice_qkv_local_cp(
1083 bias,
1084 key_dim=self.module.key_dim,
1085 value_dim=self.module.value_dim,
1086 dim=0,
1087 cp_rank=self.cp_rank,
1088 cp_size=self.cp_size,
1089 )
1091 def forward(
1092 self,
1093 hidden_states: torch.Tensor,
1094 cache_params: Optional[object] = None,
1095 attention_mask: Optional[torch.Tensor] = None,
1096 **kwargs,
1097 ) -> torch.Tensor:
1098 """Run Gated DeltaNet with pure Ulysses context parallel."""
1099 if cache_params is not None:
1100 raise NotImplementedError("GDN Ulysses CP currently supports training without cache only.")
1101 if (
1102 kwargs.get("cu_seq_lens_q") is not None
1103 or kwargs.get("seq_idx") is not None
1104 ):
1105 raise NotImplementedError("Packed variable-length GDN Ulysses CP is not implemented.")
1106 hidden_states = _local_tensor_at_cp_boundary(hidden_states)
1107 if self.gdn_backend == "triton" and hidden_states.shape[1] % 64:
1108 raise NotImplementedError(
1109 "GDN Ulysses Triton backend requires each rank's local "
1110 f"sequence length ({hidden_states.shape[1]}) divisible by 64."
1111 )
1113 base = self.module
1114 hidden_states = _apply_local_padding_mask(hidden_states, attention_mask)
1116 bsz, local_seq_len, _ = hidden_states.shape
1117 mixed_qkv = base.in_proj_qkv(hidden_states)
1118 z = base.in_proj_z(hidden_states).reshape(
1119 bsz,
1120 local_seq_len,
1121 base.num_v_heads,
1122 base.head_v_dim,
1123 )
1124 b = base.in_proj_b(hidden_states)
1125 a = base.in_proj_a(hidden_states)
1127 q_proj, k_proj, v_proj = torch.split(
1128 mixed_qkv,
1129 [base.key_dim, base.key_dim, base.value_dim],
1130 dim=-1,
1131 )
1132 q_proj, k_proj, v_proj, b, a = self._seq_to_head_qkvba(
1133 q_proj, k_proj, v_proj, b, a
1134 )
1136 full_seq_len = q_proj.shape[1]
1137 local_key_dim = base.key_dim // self.cp_size
1138 local_value_dim = base.value_dim // self.cp_size
1139 local_num_k_heads = base.num_k_heads // self.cp_size
1140 local_num_v_heads = base.num_v_heads // self.cp_size
1141 local_conv_dim = local_key_dim * 2 + local_value_dim
1143 mixed_qkv = torch.cat((q_proj, k_proj, v_proj), dim=-1).transpose(1, 2)
1144 conv_out = F.conv1d(
1145 input=mixed_qkv,
1146 weight=self._local_conv_weight(),
1147 bias=self._local_conv_bias(),
1148 stride=base.conv1d.stride,
1149 padding=base.conv1d.padding,
1150 dilation=base.conv1d.dilation,
1151 groups=local_conv_dim,
1152 )
1153 mixed_qkv = F.silu(conv_out[:, :, :full_seq_len]).transpose(1, 2)
1155 query, key, value = torch.split(
1156 mixed_qkv,
1157 [local_key_dim, local_key_dim, local_value_dim],
1158 dim=-1,
1159 )
1160 query = query.reshape(bsz, full_seq_len, local_num_k_heads, base.head_k_dim)
1161 key = key.reshape(bsz, full_seq_len, local_num_k_heads, base.head_k_dim)
1162 value = value.reshape(bsz, full_seq_len, local_num_v_heads, base.head_v_dim)
1164 a_log = _slice_local_cp(base.A_log, 0, self.cp_rank, self.cp_size)
1165 dt_bias = _slice_local_cp(base.dt_bias, 0, self.cp_rank, self.cp_size)
1166 beta = b.sigmoid()
1167 g = -a_log.float().exp() * F.softplus(a.float() + dt_bias)
1169 kv_groups = _kv_groups(base)
1170 if kv_groups > 1:
1171 query = query.repeat_interleave(kv_groups, dim=2)
1172 key = key.repeat_interleave(kv_groups, dim=2)
1174 core_attn_out, _ = chunk_gated_delta_rule(
1175 query,
1176 key,
1177 value,
1178 g=g,
1179 beta=beta,
1180 initial_state=None,
1181 output_final_state=False,
1182 chunk_size=self.chunk_size,
1183 use_qk_l2norm_in_kernel=True,
1184 backend=self.gdn_backend,
1185 )
1187 core_attn_out = self._head_to_seq(core_attn_out)
1188 core_attn_out = core_attn_out.reshape(-1, base.head_v_dim)
1189 z_flat = z.reshape(-1, base.head_v_dim)
1190 core_attn_out = base.norm(core_attn_out, z_flat)
1191 core_attn_out = core_attn_out.reshape(bsz, local_seq_len, base.value_dim)
1192 if hasattr(base, "out_proj_input"):
1193 core_attn_out = base.out_proj_input(core_attn_out)
1194 return base.out_proj(core_attn_out)
1197class GatedDeltaNetP2PCP(nn.Module):
1198 """Sequence-sharded GDN CP with an affine-summary state wavefront."""
1200 def __init__(
1201 self,
1202 module: nn.Module,
1203 device_mesh: DeviceMesh,
1204 *,
1205 backend: str = "eager",
1206 chunk_size: int = 64,
1207 ):
1208 super().__init__()
1209 self.module = module
1210 backend = backend.lower()
1211 if backend not in _GDN_BACKENDS:
1212 raise ValueError(
1213 f"unsupported GDN backend {backend!r}; "
1214 f"expected one of {sorted(_GDN_BACKENDS)}."
1215 )
1216 if chunk_size <= 0:
1217 raise ValueError(f"chunk_size must be positive, got {chunk_size}.")
1218 if backend == "triton" and chunk_size != 64:
1219 raise NotImplementedError("GDN Triton backend requires chunk_size=64.")
1220 self.gdn_backend = backend
1221 self.chunk_size = chunk_size
1222 self.cp_mesh = _ensure_1d(device_mesh)
1223 self.cp_size = self.cp_mesh.size()
1224 self.cp_rank = self.cp_mesh.get_local_rank()
1225 self._validate_module()
1227 def _validate_module(self) -> None:
1228 """Validate the Conv1d requirements of the P2P CP path."""
1229 required = (
1230 "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a",
1231 "conv1d", "A_log", "dt_bias", "norm", "out_proj",
1232 "num_k_heads", "num_v_heads", "head_k_dim", "head_v_dim",
1233 "key_dim", "value_dim", "conv_dim",
1234 )
1235 missing = [name for name in required if not hasattr(self.module, name)]
1236 if missing:
1237 raise TypeError(
1238 "Qwen3.5 GDN layer is missing required attributes: "
1239 + ", ".join(missing)
1240 )
1241 if self.gdn_backend == "triton" and (
1242 self.module.head_k_dim != 128 or self.module.head_v_dim != 128
1243 ):
1244 raise NotImplementedError(
1245 "linear attention P2P Triton backend requires "
1246 "head_k_dim=head_v_dim=128."
1247 )
1248 if self.module.num_v_heads % self.module.num_k_heads:
1249 raise ValueError("GDN num_v_heads must be divisible by num_k_heads.")
1250 _validate_gdn_activation(self.module)
1251 conv = self.module.conv1d
1252 if conv.stride != (1,):
1253 raise ValueError(
1254 "linear attention P2P CP currently supports only conv1d stride=1."
1255 )
1256 if conv.groups != self.module.conv_dim:
1257 raise ValueError(
1258 "linear attention P2P CP expects depthwise conv1d groups=conv_dim."
1259 )
1260 if (
1261 conv.in_channels != self.module.conv_dim
1262 or conv.out_channels != self.module.conv_dim
1263 ):
1264 raise ValueError(
1265 "linear attention P2P CP expects conv1d channels to match conv_dim."
1266 )
1268 def _conv1d_with_halo(self, mixed_qkv: torch.Tensor) -> torch.Tensor:
1269 """Run local Conv1d after exchanging only the previous-rank halo."""
1270 return _causal_conv1d_with_cp_halo(
1271 mixed_qkv,
1272 self.module.conv1d,
1273 self.cp_mesh,
1274 self.cp_rank,
1275 self.cp_size,
1276 )
1278 def forward(
1279 self,
1280 hidden_states: torch.Tensor,
1281 cache_params: Optional[object] = None,
1282 attention_mask: Optional[torch.Tensor] = None,
1283 **kwargs,
1284 ) -> torch.Tensor:
1285 """Run Gated DeltaNet on local sequence shards with recurrent-state P2P."""
1286 if cache_params is not None:
1287 raise NotImplementedError("GDN P2P CP currently supports training without cache only.")
1288 if (
1289 kwargs.get("cu_seq_lens_q") is not None
1290 or kwargs.get("seq_idx") is not None
1291 ):
1292 raise NotImplementedError("Packed variable-length GDN P2P is not implemented.")
1293 hidden_states = _local_tensor_at_cp_boundary(hidden_states)
1295 base = self.module
1296 hidden_states = _apply_local_padding_mask(hidden_states, attention_mask)
1298 bsz, local_seq_len, _ = hidden_states.shape
1299 mixed_qkv = base.in_proj_qkv(hidden_states)
1300 z = base.in_proj_z(hidden_states).reshape(
1301 bsz,
1302 local_seq_len,
1303 base.num_v_heads,
1304 base.head_v_dim,
1305 )
1306 b = base.in_proj_b(hidden_states)
1307 a = base.in_proj_a(hidden_states)
1309 mixed_qkv = self._conv1d_with_halo(mixed_qkv)
1310 query, key, value = torch.split(
1311 mixed_qkv,
1312 [base.key_dim, base.key_dim, base.value_dim],
1313 dim=-1,
1314 )
1315 query = query.reshape(bsz, local_seq_len, base.num_k_heads, base.head_k_dim)
1316 key = key.reshape(bsz, local_seq_len, base.num_k_heads, base.head_k_dim)
1317 value = value.reshape(bsz, local_seq_len, base.num_v_heads, base.head_v_dim)
1319 beta = b.sigmoid()
1320 g = -base.A_log.float().exp() * F.softplus(a.float() + base.dt_bias)
1322 kv_groups = _kv_groups(base)
1323 if kv_groups > 1:
1324 query = query.repeat_interleave(kv_groups, dim=2)
1325 key = key.repeat_interleave(kv_groups, dim=2)
1327 if self.gdn_backend == "triton":
1328 if local_seq_len % 64 != 0:
1329 raise NotImplementedError(
1330 "linear attention P2P Triton backend requires each CP "
1331 f"rank's local sequence length ({local_seq_len}) to be "
1332 "divisible by 64."
1333 )
1334 if not is_triton_gdn_available(query, key, value, g, beta):
1335 raise RuntimeError(
1336 "linear attention P2P Triton backend requires an NPU "
1337 "input satisfying the fixed GDN contract and a validated "
1338 "triton-ascend >= 3.2.1 installation."
1339 )
1340 core_attn_out = _gdn_state_p2p_triton(
1341 query,
1342 key,
1343 value,
1344 g,
1345 beta,
1346 self.cp_mesh,
1347 self.cp_rank,
1348 self.cp_size,
1349 )
1350 else:
1351 core_attn_out = _gdn_state_p2p_summary(
1352 query,
1353 key,
1354 value,
1355 g,
1356 beta,
1357 self.cp_mesh,
1358 self.cp_rank,
1359 self.cp_size,
1360 chunk_size=self.chunk_size,
1361 use_qk_l2norm_in_kernel=True,
1362 )
1364 core_attn_out = core_attn_out.reshape(-1, base.head_v_dim)
1365 z_flat = z.reshape(-1, base.head_v_dim)
1366 core_attn_out = base.norm(core_attn_out, z_flat)
1367 core_attn_out = core_attn_out.reshape(bsz, local_seq_len, base.value_dim)
1368 if hasattr(base, "out_proj_input"):
1369 core_attn_out = base.out_proj_input(core_attn_out)
1370 return base.out_proj(core_attn_out)