Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / gated_delta_net.py: 26%
762 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07: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."""
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"]
25import functools
26import inspect
27from typing import Any, Callable, NamedTuple, Optional
29import torch
31import torch.distributed as dist
32from torch import nn
33from torch.nn import functional as F
34from torch.utils.checkpoint import checkpoint
36from hyper_parallel.components.functional.linear_attention_convolution import (
37 fused_causal_conv1d,
38)
39from hyper_parallel.components.functional.linear_attention_packed import (
40 PackedSequenceInfo, packed_causal_conv1d, resolve_packed_sequence,
41)
43from hyper_parallel.core.context_parallel.context_parallel import (
44 _ensure_1d,
45)
46from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
47from hyper_parallel.core.dtensor.dtensor import DTensor
48from hyper_parallel.components.modules.gated_delta_net import (
49 chunk_gated_delta_rule,
50 is_triton_gdn_available,
51 torch_chunk_gated_delta_rule,
52)
53from hyper_parallel.distributed.context_parallel.collectives import ulysses_head_to_seq, ulysses_seq_to_head
54from hyper_parallel.core.utils import communication
57_GDN_BACKENDS = frozenset({"eager", "triton", "ascendc"})
60def _global_peer_rank(cp_mesh: DeviceMesh, local_rank: int) -> int:
61 """Map a CP-local rank index to its global distributed rank."""
62 return int(cp_mesh.rank_list[local_rank])
65def _slice_local_cp(
66 tensor: torch.Tensor,
67 dim: int,
68 cp_rank: int,
69 cp_size: int,
70) -> torch.Tensor:
71 """Return this CP rank's contiguous slice along ``dim``."""
72 dim_size = tensor.shape[dim]
73 if dim_size % cp_size != 0:
74 raise ValueError(
75 f"linear attention CP expects dim size {dim_size} "
76 f"to be divisible by cp_size {cp_size}."
77 )
78 chunk = dim_size // cp_size
79 return tensor.narrow(dim, cp_rank * chunk, chunk)
82def _slice_qkv_local_cp(
83 tensor: torch.Tensor,
84 *,
85 key_dim: int,
86 value_dim: int,
87 dim: int,
88 cp_rank: int,
89 cp_size: int,
90) -> torch.Tensor:
91 """Slice a fused ``[Q, K, V]`` tensor on the Q/K/V channel dimension."""
92 q, k, v = torch.split(tensor, [key_dim, key_dim, value_dim], dim=dim)
93 return torch.cat(
94 (
95 _slice_local_cp(q, dim, cp_rank, cp_size),
96 _slice_local_cp(k, dim, cp_rank, cp_size),
97 _slice_local_cp(v, dim, cp_rank, cp_size),
98 ),
99 dim=dim,
100 )
103def _local_tensor_at_cp_boundary(tensor: torch.Tensor) -> torch.Tensor:
104 """Return the local tensor carried by a CP-boundary input.
106 The first supported Qwen3.5 linear-attention CP path keeps decoder-layer
107 activations as local sequence shards. If an upstream wrapper passes that
108 shard as a DTensor, use its local tensor and continue with the same
109 ``[B, S_local, H]`` boundary contract.
110 """
111 if isinstance(tensor, DTensor):
112 return tensor.to_local()
113 return tensor
116def _kv_groups(module: nn.Module) -> int:
117 """Return the Q/K replication factor used by the official HF GDN layer."""
118 return int(module.num_v_heads // module.num_k_heads)
121def _apply_local_padding_mask(
122 hidden_states: torch.Tensor,
123 attention_mask: Optional[torch.Tensor],
124) -> torch.Tensor:
125 """Apply the supported local 2-D padding mask without silent fallback."""
126 if attention_mask is None:
127 return hidden_states
128 if attention_mask.ndim != 2:
129 raise NotImplementedError(
130 "Qwen3.5 GDN Context Parallel supports only a local 2-D padding mask."
131 )
132 if tuple(attention_mask.shape) != tuple(hidden_states.shape[:2]):
133 raise ValueError(
134 "Qwen3.5 GDN Context Parallel expects attention_mask shape "
135 f"{tuple(hidden_states.shape[:2])}, got {tuple(attention_mask.shape)}."
136 )
137 return hidden_states * attention_mask[:, :, None].to(hidden_states.dtype)
140def _validate_gdn_activation(module: nn.Module) -> None:
141 """Reject model variants whose Conv activation is not implemented here."""
142 activation = getattr(module, "activation", "silu")
143 if activation not in ("silu", "swish"):
144 raise NotImplementedError(
145 "Qwen3.5 GDN Context Parallel currently supports only SiLU/Swish "
146 f"Conv activation, got {activation!r}."
147 )
150def _all_to_all_previous_rank_halo(
151 tail: torch.Tensor,
152 cp_mesh: DeviceMesh,
153 cp_rank: int,
154 cp_size: int,
155) -> torch.Tensor:
156 """Send a convolution halo only to the next rank using differentiable A2AV."""
157 if cp_size == 1:
158 return torch.zeros_like(tail)
160 cp_group = cp_mesh.get_group()
161 group_ranks = tuple(int(rank) for rank in dist.get_process_group_ranks(cp_group))
162 rank_list = tuple(int(rank) for rank in cp_mesh.rank_list)
163 rank_to_group_index = {rank: index for index, rank in enumerate(group_ranks)}
164 halo_width = tail.shape[1]
166 input_splits = [0] * cp_size
167 exchange_input = tail.permute(1, 0, 2).contiguous()
168 if cp_rank < cp_size - 1:
169 input_splits[rank_to_group_index[rank_list[cp_rank + 1]]] = halo_width
170 else:
171 exchange_input = exchange_input[:0]
173 output_splits = [0] * cp_size
174 if cp_rank > 0:
175 output_splits[rank_to_group_index[rank_list[cp_rank - 1]]] = halo_width
177 exchange_output = communication.differentiable_all_to_all_single(
178 exchange_input,
179 input_splits,
180 output_splits,
181 group=cp_group,
182 )
183 if cp_rank == 0:
184 return torch.zeros_like(tail) + exchange_output.sum().to(tail.dtype) * 0
185 return exchange_output.permute(1, 0, 2).contiguous()
188def _causal_conv1d_with_cp_halo(
189 mixed_qkv: torch.Tensor,
190 conv1d: nn.Conv1d,
191 cp_mesh: DeviceMesh,
192 cp_rank: int,
193 cp_size: int,
194 sequence_info: Optional[PackedSequenceInfo] = None,
195) -> torch.Tensor:
196 """Run causal depthwise Conv1d with only the previous rank's boundary."""
197 kernel_size = conv1d.kernel_size[0]
198 dilation = conv1d.dilation[0]
199 halo_width = (kernel_size - 1) * dilation
200 if sequence_info is not None:
201 halo = None
202 if halo_width and cp_size > 1:
203 if mixed_qkv.shape[1] < halo_width:
204 raise ValueError("Packed CP requires a local interval at least as long as its convolution halo.")
205 halo = _all_to_all_previous_rank_halo(
206 mixed_qkv[:, -halo_width:].contiguous(), cp_mesh, cp_rank, cp_size)
207 return packed_causal_conv1d(mixed_qkv, conv1d, sequence_info, halo)
208 if halo_width == 0 or cp_size == 1:
209 conv_out = conv1d(mixed_qkv.transpose(1, 2))
210 return F.silu(conv_out[:, :, : mixed_qkv.shape[1]]).transpose(1, 2)
212 if mixed_qkv.shape[1] < halo_width:
213 raise ValueError(
214 "linear attention CP conv halo requires local_seq_len >= "
215 f"{halo_width}, got {mixed_qkv.shape[1]}."
216 )
218 halo = _all_to_all_previous_rank_halo(
219 mixed_qkv[:, -halo_width:, :].contiguous(),
220 cp_mesh,
221 cp_rank,
222 cp_size,
223 )
224 conv_input = torch.cat((halo, mixed_qkv), dim=1).transpose(1, 2)
225 conv_out = F.conv1d(
226 input=conv_input,
227 weight=conv1d.weight,
228 bias=conv1d.bias,
229 stride=conv1d.stride,
230 padding=0,
231 dilation=conv1d.dilation,
232 groups=conv1d.groups,
233 )
234 return F.silu(conv_out).transpose(1, 2)
237def _l2norm_torch(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor:
238 """Match the pure torch GDN reference l2norm helper."""
239 return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
242class _GDNPreparedChunks(NamedTuple):
243 """Reusable chunk intermediates shared by state-summary CP modes."""
245 initial_dtype: torch.dtype
246 query: torch.Tensor
247 key: torch.Tensor
248 chunk_value: torch.Tensor
249 g: torch.Tensor
250 decay_mask: torch.Tensor
251 k_cumdecay: torch.Tensor
252 sequence_length: int
253 total_sequence_length: int
254 chunk_size: int
257def _prepare_gdn_chunks_for_summary(
258 query: torch.Tensor,
259 key: torch.Tensor,
260 value: torch.Tensor,
261 g: torch.Tensor,
262 beta: torch.Tensor,
263 *,
264 chunk_size: int = 64,
265 use_qk_l2norm_in_kernel: bool = False,
266) -> _GDNPreparedChunks:
267 """Prepare GDN chunk intermediates shared by summary and local output."""
268 initial_dtype = query.dtype
269 if use_qk_l2norm_in_kernel:
270 query = _l2norm_torch(query, dim=-1, eps=1e-6)
271 key = _l2norm_torch(key, dim=-1, eps=1e-6)
273 query, key, value, beta, g = [
274 x.transpose(1, 2).contiguous().to(torch.float32)
275 for x in (query, key, value, beta, g)
276 ]
278 sequence_length = key.shape[2]
279 pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
280 query = F.pad(query, (0, 0, 0, pad_size))
281 key = F.pad(key, (0, 0, 0, pad_size))
282 value = F.pad(value, (0, 0, 0, pad_size))
283 beta = F.pad(beta, (0, pad_size))
284 g = F.pad(g, (0, pad_size))
285 total_sequence_length = sequence_length + pad_size
287 query = query * (1 / (query.shape[-1] ** 0.5))
288 v_beta = value * beta.unsqueeze(-1)
289 k_beta = key * beta.unsqueeze(-1)
290 query, key, k_beta, v_beta = [
291 x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
292 for x in (query, key, k_beta, v_beta)
293 ]
294 g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
296 mask = torch.triu(
297 torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
298 diagonal=0,
299 )
300 g = g.cumsum(dim=-1)
301 decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
302 attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
303 for row_idx in range(1, chunk_size):
304 row = attn[..., row_idx, :row_idx].clone()
305 sub = attn[..., :row_idx, :row_idx].clone()
306 attn[..., row_idx, :row_idx] = row + (row.unsqueeze(-1) * sub).sum(-2)
307 attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
309 chunk_value = attn @ v_beta
310 k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
311 return _GDNPreparedChunks(
312 initial_dtype=initial_dtype,
313 query=query,
314 key=key,
315 chunk_value=chunk_value,
316 g=g,
317 decay_mask=decay_mask,
318 k_cumdecay=k_cumdecay,
319 sequence_length=sequence_length,
320 total_sequence_length=total_sequence_length,
321 chunk_size=chunk_size,
322 )
325def _compute_gdn_state_summary_from_prepared(
326 prepared: _GDNPreparedChunks,
327) -> tuple[torch.Tensor, torch.Tensor]:
328 """Compute ``state_out = M @ state_in + S`` from prepared GDN chunks."""
329 key = prepared.key
330 batch_size, num_heads, _, _, k_head_dim = key.shape
331 v_head_dim = prepared.chunk_value.shape[-1]
332 eye = torch.eye(k_head_dim, device=key.device, dtype=torch.float32).reshape(
333 1, 1, k_head_dim, k_head_dim
334 )
335 state_ext = torch.zeros(
336 batch_size,
337 num_heads,
338 k_head_dim,
339 v_head_dim,
340 device=key.device,
341 dtype=torch.float32,
342 )
343 transition = eye.expand(batch_size, num_heads, -1, -1).clone()
345 for chunk_idx in range(key.shape[2]):
346 key_i = key[:, :, chunk_idx]
347 value_i = prepared.chunk_value[:, :, chunk_idx]
348 w_i = prepared.k_cumdecay[:, :, chunk_idx]
349 g_i = prepared.g[:, :, chunk_idx]
350 decay = g_i[:, :, -1].exp()
351 key_decay = key_i * (g_i[:, :, -1, None] - g_i).exp()[..., None]
353 transition_i = (
354 decay[:, :, None, None] * eye
355 - key_decay.transpose(-1, -2) @ w_i
356 )
357 state_ext_i = key_decay.transpose(-1, -2) @ value_i
358 state_ext = transition_i @ state_ext + state_ext_i
359 transition = transition_i @ transition
361 return state_ext, transition
364def _checkpoint_gdn_state_summary(
365 prepared: _GDNPreparedChunks,
366) -> tuple[torch.Tensor, torch.Tensor]:
367 """Compute a state summary without retaining its per-chunk autograd graph."""
368 if not torch.is_grad_enabled():
369 return _compute_gdn_state_summary_from_prepared(prepared)
371 def recompute(
372 key: torch.Tensor,
373 chunk_value: torch.Tensor,
374 g: torch.Tensor,
375 k_cumdecay: torch.Tensor,
376 ) -> tuple[torch.Tensor, torch.Tensor]:
377 """Rebuild a prepared view from explicit checkpoint inputs."""
378 checkpoint_prepared = prepared._replace(
379 key=key,
380 chunk_value=chunk_value,
381 g=g,
382 k_cumdecay=k_cumdecay,
383 )
384 return _compute_gdn_state_summary_from_prepared(checkpoint_prepared)
386 return checkpoint(
387 recompute,
388 prepared.key,
389 prepared.chunk_value,
390 prepared.g,
391 prepared.k_cumdecay,
392 use_reentrant=False,
393 preserve_rng_state=False,
394 )
397def _run_prepared_gdn_chunks(
398 prepared: _GDNPreparedChunks,
399 initial_state: Optional[torch.Tensor],
400) -> torch.Tensor:
401 """Run local GDN output using already prepared chunk intermediates."""
402 query = prepared.query
403 key = prepared.key
404 chunk_value = prepared.chunk_value
405 batch_size, num_heads, _, _, k_head_dim = key.shape
406 v_head_dim = chunk_value.shape[-1]
407 recurrent_state = (
408 torch.zeros(
409 batch_size,
410 num_heads,
411 k_head_dim,
412 v_head_dim,
413 device=chunk_value.device,
414 dtype=chunk_value.dtype,
415 )
416 if initial_state is None
417 else initial_state.to(chunk_value)
418 )
419 core_attn_out = torch.zeros_like(chunk_value)
421 for chunk_idx in range(0, prepared.total_sequence_length // prepared.chunk_size):
422 q_i = query[:, :, chunk_idx]
423 k_i = key[:, :, chunk_idx]
424 v_i = chunk_value[:, :, chunk_idx]
425 attn = q_i @ k_i.transpose(-1, -2) * prepared.decay_mask[:, :, chunk_idx]
426 v_prime = prepared.k_cumdecay[:, :, chunk_idx] @ recurrent_state
427 v_new = v_i - v_prime
428 attn_inter = (
429 q_i * prepared.g[:, :, chunk_idx, :, None].exp()
430 ) @ recurrent_state
431 core_attn_out[:, :, chunk_idx] = attn_inter + attn @ v_new
432 recurrent_state = (
433 recurrent_state * prepared.g[:, :, chunk_idx, -1, None, None].exp()
434 + (
435 k_i
436 * (
437 prepared.g[:, :, chunk_idx, -1, None]
438 - prepared.g[:, :, chunk_idx]
439 ).exp()[..., None]
440 ).transpose(-1, -2) @ v_new
441 )
443 core_attn_out = core_attn_out.reshape(
444 core_attn_out.shape[0],
445 core_attn_out.shape[1],
446 -1,
447 core_attn_out.shape[-1],
448 )
449 core_attn_out = core_attn_out[:, :, :prepared.sequence_length]
450 return core_attn_out.transpose(1, 2).contiguous().to(prepared.initial_dtype)
453class _RecvInitialStateP2PFunction(torch.autograd.Function):
454 """Receive the recurrent initial state; send its gradient in backward."""
456 @staticmethod
457 def forward( # pylint: disable=arguments-differ
458 ctx,
459 anchor: torch.Tensor,
460 cp_group,
461 prev_rank: int,
462 state_shape: tuple[int, ...],
463 ) -> torch.Tensor:
464 """Receive the initial state from the preceding CP rank."""
465 state = torch.empty(state_shape, device=anchor.device, dtype=torch.float32)
466 dist.recv(state, src=prev_rank, group=cp_group)
467 ctx.cp_group = cp_group
468 ctx.prev_rank = prev_rank
469 return state
471 @staticmethod
472 def backward(ctx, grad_state: Optional[torch.Tensor]):
473 """Return the incoming state gradient to the preceding CP rank."""
474 if grad_state is None:
475 raise RuntimeError("linear attention P2P backward missing initial-state grad.")
476 dist.send(grad_state.contiguous(), dst=ctx.prev_rank, group=ctx.cp_group)
477 return None, None, None, None
480class _SendFinalStateP2PFunction(torch.autograd.Function):
481 """Send the recurrent final state; receive its gradient in backward."""
483 @staticmethod
484 def forward( # pylint: disable=arguments-differ
485 ctx,
486 final_state: torch.Tensor,
487 cp_group,
488 next_rank: int,
489 ) -> torch.Tensor:
490 """Send the final state to the succeeding CP rank."""
491 dist.send(final_state.contiguous(), dst=next_rank, group=cp_group)
492 ctx.cp_group = cp_group
493 ctx.next_rank = next_rank
494 ctx.state_shape = tuple(final_state.shape)
495 ctx.state_dtype = final_state.dtype
496 return final_state.new_zeros(())
498 @staticmethod
499 def backward(ctx, grad_token: torch.Tensor):
500 """Receive the final-state gradient from the succeeding CP rank."""
501 grad_state = torch.empty(
502 ctx.state_shape,
503 device=grad_token.device,
504 dtype=ctx.state_dtype,
505 )
506 dist.recv(grad_state, src=ctx.next_rank, group=ctx.cp_group)
507 return grad_state, None, None
510def _apply_gdn_state_summary(
511 state_ext: torch.Tensor,
512 transition: torch.Tensor,
513 initial_state: Optional[torch.Tensor],
514) -> torch.Tensor:
515 """Apply ``state_out = M @ state_in + S`` to an incoming GDN state."""
516 if initial_state is None:
517 return state_ext
518 return transition @ initial_state.to(transition) + state_ext
521def _gdn_state_p2p_summary(
522 query: torch.Tensor,
523 key: torch.Tensor,
524 value: torch.Tensor,
525 g: torch.Tensor,
526 beta: torch.Tensor,
527 cp_mesh: DeviceMesh,
528 cp_rank: int,
529 cp_size: int,
530 *,
531 chunk_size: int,
532 use_qk_l2norm_in_kernel: bool,
533 sequence_info: Optional[PackedSequenceInfo] = None,
534) -> torch.Tensor:
535 """Run local GDN with an affine-summary state wavefront.
537 Every rank prepares its local chunks and state transition in parallel.
538 The rank-ordered critical path then contains only ``M @ state + S`` and
539 the small state transfer. Token outputs retain the ordinary PyTorch graph,
540 while the two custom autograd boundaries reverse the state communication.
541 """
542 if sequence_info is not None:
543 cp_group = None if cp_mesh is None else cp_mesh.get_group()
544 incoming = None
545 if cp_rank > 0:
546 incoming = _RecvInitialStateP2PFunction.apply(
547 query, cp_group, _global_peer_rank(cp_mesh, cp_rank - 1),
548 (1, value.shape[2], query.shape[-1], value.shape[-1]))
549 output, finals = torch_chunk_gated_delta_rule(
550 query, key, value, g, beta, chunk_size=chunk_size,
551 initial_state=sequence_info.first_state(incoming), output_final_state=True,
552 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, cu_seqlens=sequence_info.local_offsets)
553 if incoming is not None:
554 output = output + incoming.sum().to(output.dtype) * 0
555 if cp_rank < cp_size - 1:
556 token = _SendFinalStateP2PFunction.apply(
557 finals[-1:].contiguous(), cp_group, _global_peer_rank(cp_mesh, cp_rank + 1))
558 output = output + token.to(output.dtype) * 0
559 return output
560 if cp_size == 1:
561 core_attn_out, _ = torch_chunk_gated_delta_rule(
562 query,
563 key,
564 value,
565 g=g,
566 beta=beta,
567 initial_state=None,
568 output_final_state=False,
569 chunk_size=chunk_size,
570 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
571 )
572 return core_attn_out
574 cp_group = cp_mesh.get_group()
575 prev_rank = _global_peer_rank(cp_mesh, cp_rank - 1) if cp_rank > 0 else -1
576 next_rank = _global_peer_rank(cp_mesh, cp_rank + 1) if cp_rank < cp_size - 1 else -1
577 state_shape = (query.shape[0], value.shape[2], query.shape[3], value.shape[3])
579 prepared = _prepare_gdn_chunks_for_summary(
580 query,
581 key,
582 value,
583 g,
584 beta,
585 chunk_size=chunk_size,
586 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
587 )
588 state_ext = None
589 transition = None
590 if cp_rank < cp_size - 1:
591 state_ext, transition = _checkpoint_gdn_state_summary(prepared)
593 initial_state = None
594 if cp_rank > 0:
595 initial_state = _RecvInitialStateP2PFunction.apply(
596 query,
597 cp_group,
598 prev_rank,
599 state_shape,
600 )
602 send_token = None
603 if cp_rank < cp_size - 1:
604 final_state = _apply_gdn_state_summary(
605 state_ext,
606 transition,
607 initial_state,
608 )
609 send_token = _SendFinalStateP2PFunction.apply(final_state, cp_group, next_rank)
611 core_attn_out = _run_prepared_gdn_chunks(prepared, initial_state)
612 if send_token is not None:
613 core_attn_out = core_attn_out + send_token.to(core_attn_out.dtype) * 0
615 return core_attn_out
618class _GDNStateP2PTritonFunction(torch.autograd.Function):
619 """Pipeline fused affine GDN states over sequence-sharded CP ranks."""
621 @staticmethod
622 def forward( # pylint: disable=arguments-differ,too-many-locals
623 ctx,
624 query: torch.Tensor,
625 key: torch.Tensor,
626 value: torch.Tensor,
627 g: torch.Tensor,
628 beta: torch.Tensor,
629 cp_rank: int,
630 cp_size: int,
631 cp_group,
632 prev_rank: int,
633 next_rank: int,
634 sequence_info: Optional[PackedSequenceInfo],
635 ) -> torch.Tensor:
636 """Run fused local GDN and forward its affine state across CP ranks."""
637 from hyper_parallel.components.functional.gated_delta_net import ( # pylint: disable=import-outside-toplevel
638 chunk_gated_delta_rule_fwd_apply_state_saved,
639 chunk_gated_delta_rule_fwd_output_saved,
640 chunk_gated_delta_rule_fwd_prepare_saved,
641 )
642 from hyper_parallel.components.functional.gated_delta_net_state_summary import ( # pylint: disable=import-outside-toplevel
643 apply_gdn_state_summary,
644 chunk_gated_delta_rule_state_summary_fwd,
645 )
647 receive_previous = cp_rank > 0 and (sequence_info is None or sequence_info.continues_previous)
648 send_next = cp_rank < cp_size - 1 and (sequence_info is None or sequence_info.continues_next)
649 initial_state = None
650 recv_buffer = None
651 recv_handle = None
652 if receive_previous:
653 recv_buffer = torch.empty(
654 (query.shape[0], query.shape[2], query.shape[3], value.shape[3]),
655 device=query.device,
656 dtype=torch.float32,
657 )
658 recv_handle = dist.irecv(recv_buffer, src=prev_rank, group=cp_group)
660 cu_seqlens = None if sequence_info is None else sequence_info.materialize(query.device, 64).cu_seqlens
661 (
662 query_norm,
663 key_norm,
664 _,
665 _,
666 g_cumsum,
667 matrix_a,
668 w,
669 u,
670 scale,
671 ) = chunk_gated_delta_rule_fwd_prepare_saved(
672 query,
673 key,
674 value,
675 g,
676 beta,
677 use_qk_l2norm_in_kernel=False,
678 cu_seqlens=cu_seqlens,
679 )
681 state_ext = None
682 transition = None
683 if send_next:
684 tail = slice(None) if sequence_info is None else sequence_info.tail
685 state_ext, transition = chunk_gated_delta_rule_state_summary_fwd(
686 key_norm[:, tail], w[:, tail], u[:, tail], g_cumsum[:, tail],
687 output_transition=sequence_info is None or sequence_info.transfer_active,
688 )
690 if recv_handle is not None and (sequence_info is None or sequence_info.transfer_active):
691 recv_handle.wait()
692 initial_state = recv_buffer
693 recv_handle = None
695 send_buffer = None
696 send_handle = None
697 if send_next:
698 send_buffer = apply_gdn_state_summary(
699 state_ext,
700 transition,
701 initial_state if transition is not None else None,
702 ).contiguous()
703 send_handle = dist.isend(send_buffer, dst=next_rank, group=cp_group)
705 if recv_handle is not None:
706 recv_handle.wait()
707 initial_state = recv_buffer
708 h, v_new, _ = chunk_gated_delta_rule_fwd_apply_state_saved(
709 key_norm,
710 g_cumsum,
711 w,
712 u,
713 initial_state=initial_state if sequence_info is None else sequence_info.first_state(initial_state),
714 output_final_state=False,
715 cu_seqlens=cu_seqlens,
716 )
717 output = chunk_gated_delta_rule_fwd_output_saved(
718 query_norm,
719 key_norm,
720 g_cumsum,
721 h,
722 v_new,
723 scale,
724 cu_seqlens=cu_seqlens,
725 ).to(query.dtype)
727 if send_handle is not None:
728 send_handle.wait()
729 del send_buffer
731 empty = query.new_empty(0)
732 ctx.save_for_backward(
733 query_norm,
734 key_norm,
735 value,
736 g_cumsum,
737 beta,
738 matrix_a,
739 initial_state if initial_state is not None else empty,
740 transition if transition is not None else empty,
741 cu_seqlens if cu_seqlens is not None else query.new_empty(0, dtype=torch.int64),
742 )
743 ctx.sequence_info = sequence_info
744 ctx.has_initial_state = initial_state is not None
745 ctx.has_transition = transition is not None
746 ctx.receive_previous = receive_previous
747 ctx.send_next = send_next
748 ctx.cp_rank = cp_rank
749 ctx.cp_size = cp_size
750 ctx.cp_group = cp_group
751 ctx.prev_rank = prev_rank
752 ctx.next_rank = next_rank
753 ctx.scale = scale
754 return output
756 @staticmethod
757 def backward(ctx, grad_output: torch.Tensor): # pylint: disable=too-many-locals
758 """Backpropagate local GDN tensors and the state gradient wavefront."""
759 from hyper_parallel.components.functional.gated_delta_net import ( # pylint: disable=import-outside-toplevel
760 chunk_gated_delta_rule_bwd_finish_saved,
761 chunk_gated_delta_rule_bwd_prepare_saved,
762 chunk_gated_delta_rule_bwd_state_saved,
763 )
764 from hyper_parallel.components.functional.gated_delta_net_state_summary import ( # pylint: disable=import-outside-toplevel
765 apply_gdn_state_gradient_summary,
766 chunk_gated_delta_rule_state_gradient_summary_bwd,
767 )
769 (
770 query,
771 key,
772 value,
773 g_cumsum,
774 beta,
775 matrix_a,
776 initial_state,
777 transition,
778 saved_cu_seqlens,
779 ) = ctx.saved_tensors
780 sequence_info = ctx.sequence_info
781 cu_seqlens = None if sequence_info is None else saved_cu_seqlens
782 if not ctx.has_initial_state:
783 initial_state = None
785 grad_final_state = None
786 recv_handle = None
787 if ctx.send_next:
788 recv_buffer = torch.empty(
789 (query.shape[0], query.shape[2], query.shape[3], value.shape[3]),
790 device=grad_output.device,
791 dtype=torch.float32,
792 )
793 recv_handle = dist.irecv(
794 recv_buffer,
795 src=ctx.next_rank,
796 group=ctx.cp_group,
797 )
799 kernel_initial = initial_state if sequence_info is None else sequence_info.first_state(initial_state)
800 w, h, v_new, dv = chunk_gated_delta_rule_bwd_prepare_saved(
801 query,
802 key,
803 value,
804 g_cumsum,
805 beta,
806 matrix_a,
807 kernel_initial,
808 grad_output,
809 ctx.scale,
810 cu_seqlens=cu_seqlens,
811 )
813 grad_state_ext = None
814 if ctx.receive_previous:
815 head = slice(None) if sequence_info is None else sequence_info.head
816 grad_state_ext = chunk_gated_delta_rule_state_gradient_summary_bwd(
817 query[:, head], key[:, head], w[:, head], g_cumsum[:, head],
818 grad_output[:, head], dv[:, head], ctx.scale)
820 if recv_handle is not None and (sequence_info is None or sequence_info.transfer_active):
821 recv_handle.wait()
822 grad_final_state = recv_buffer
823 recv_handle = None
825 send_buffer = None
826 send_handle = None
827 if ctx.receive_previous:
828 send_buffer = apply_gdn_state_gradient_summary(
829 grad_state_ext,
830 transition,
831 grad_final_state if ctx.has_transition else None,
832 ).contiguous()
833 send_handle = dist.isend(
834 send_buffer,
835 dst=ctx.prev_rank,
836 group=ctx.cp_group,
837 )
839 if recv_handle is not None:
840 recv_handle.wait()
841 grad_final_state = recv_buffer
842 kernel_grad_final = grad_final_state if sequence_info is None else sequence_info.last_state(grad_final_state)
843 dh, _, dv = chunk_gated_delta_rule_bwd_state_saved(
844 query,
845 key,
846 g_cumsum,
847 w,
848 kernel_initial,
849 kernel_grad_final,
850 grad_output,
851 dv,
852 ctx.scale,
853 cu_seqlens=cu_seqlens,
854 )
855 empty = query.new_empty(0)
856 dq, dk, dv, dg, dbeta = chunk_gated_delta_rule_bwd_finish_saved(
857 query,
858 key,
859 query,
860 key,
861 value,
862 g_cumsum,
863 beta,
864 matrix_a,
865 w,
866 h,
867 v_new,
868 dv,
869 grad_output,
870 dh,
871 empty,
872 empty,
873 ctx.scale,
874 use_qk_l2norm_in_kernel=False,
875 cu_seqlens=cu_seqlens,
876 )
878 if send_handle is not None:
879 send_handle.wait()
880 return dq, dk, dv, dg, dbeta, None, None, None, None, None, None
883def _gdn_state_p2p_triton(
884 query: torch.Tensor,
885 key: torch.Tensor,
886 value: torch.Tensor,
887 g: torch.Tensor,
888 beta: torch.Tensor,
889 cp_mesh: DeviceMesh,
890 cp_rank: int,
891 cp_size: int,
892 sequence_info: Optional[PackedSequenceInfo] = None,
893) -> torch.Tensor:
894 """Run fused local GDN with an affine state wavefront."""
895 if cp_size == 1:
896 output, _ = chunk_gated_delta_rule(
897 query,
898 key,
899 value,
900 g=g,
901 beta=beta,
902 output_final_state=False,
903 use_qk_l2norm_in_kernel=True,
904 backend="triton",
905 cu_seqlens=None if sequence_info is None else sequence_info.materialize(query.device, 64).cu_seqlens,
906 )
907 return output
909 query = _l2norm_torch(query)
910 key = _l2norm_torch(key)
911 prev_rank = _global_peer_rank(cp_mesh, cp_rank - 1) if cp_rank > 0 else -1
912 next_rank = (
913 _global_peer_rank(cp_mesh, cp_rank + 1) if cp_rank < cp_size - 1 else -1
914 )
915 return _GDNStateP2PTritonFunction.apply(
916 query,
917 key,
918 value,
919 g,
920 beta,
921 cp_rank,
922 cp_size,
923 cp_mesh.get_group(),
924 prev_rank,
925 next_rank,
926 sequence_info,
927 )
930def _differentiable_all_to_all_shard(
931 tensor: torch.Tensor,
932 device_mesh: DeviceMesh,
933 *,
934 split_dim: int,
935 concat_dim: int,
936) -> torch.Tensor:
937 """Split local data on ``split_dim`` and concatenate peers on ``concat_dim``.
939 This is the local-tensor equivalent of DTensor ``Shard(concat_dim) ->
940 Shard(split_dim)`` redistribution for a 1-D mesh. It uses a
941 differentiable all-to-all directly to avoid wrapping each activation in a
942 temporary DTensor.
943 """
944 split_count = device_mesh.size()
945 if split_count == 1:
946 return tensor
948 original_shape = tuple(tensor.shape)
949 dim_size = original_shape[split_dim]
950 if dim_size % split_count != 0:
951 raise ValueError(
952 f"linear attention all-to-all split dim {split_dim} with size "
953 f"{dim_size} must be divisible by cp_size {split_count}."
954 )
956 split_size = dim_size // split_count
957 final_shape = list(original_shape)
958 if split_dim != concat_dim:
959 final_shape[split_dim] = split_size
960 final_shape[concat_dim] = final_shape[concat_dim] * split_count
961 final_shape = tuple(final_shape)
963 reshape_dims = list(original_shape)
964 reshape_dims[split_dim] = split_count
965 reshape_dims.insert(split_dim + 1, split_size)
967 trans_dims = list(range(len(reshape_dims)))
968 trans_dims.remove(split_dim)
969 trans_dims.insert(0, split_dim)
971 a2a_input = tensor.reshape(reshape_dims).permute(trans_dims).contiguous()
972 reshape_shape = list(a2a_input.shape)
973 reshape_shape[0] = reshape_shape[0] * reshape_shape[1]
974 reshape_shape.pop(1)
975 a2a_input = a2a_input.reshape(reshape_shape)
977 a2a_input = a2a_input.contiguous()
978 split_len = a2a_input.shape[0] // split_count
979 input_splits = [split_len] * split_count
980 output_splits = [split_len] * split_count
981 output = communication.differentiable_all_to_all_single(
982 a2a_input,
983 input_splits,
984 output_splits,
985 group=device_mesh.get_group(),
986 )
988 output_reshape = list(output.shape)
989 output_reshape[0] = split_count
990 output_reshape.insert(1, output.shape[0] // split_count)
992 out_trans_dims = list(range(len(output_reshape)))
993 first_dim = out_trans_dims.pop(0)
994 if concat_dim >= len(out_trans_dims):
995 out_trans_dims.append(first_dim)
996 else:
997 out_trans_dims.insert(concat_dim, first_dim)
999 final_output = output.reshape(output_reshape).permute(out_trans_dims).contiguous()
1000 final_reshape = list(final_output.shape)
1001 if concat_dim < len(final_reshape) - 1:
1002 final_reshape[concat_dim] = (
1003 final_reshape[concat_dim] * final_reshape[concat_dim + 1]
1004 )
1005 final_reshape.pop(concat_dim + 1)
1007 return final_output.reshape(final_reshape).view(final_shape)
1010class GatedDeltaNetUlyssesCP(nn.Module):
1011 """Pure-Ulysses CP execution wrapper for a Qwen3.5 Gated DeltaNet module.
1013 Parameters stay owned by the original module. The wrapper only changes the
1014 execution layout:
1016 ``[B, S_local, full_heads] -> [B, S_full, local_heads] ->
1017 [B, S_local, full_heads]``.
1019 ``backend`` selects both the local attention core and convolution for dense
1020 and packed inputs.
1021 """
1023 def __init__(
1024 self,
1025 module: nn.Module,
1026 device_mesh: Optional[DeviceMesh],
1027 *,
1028 backend: str = "eager",
1029 chunk_size: int = 64,
1030 ) -> None:
1031 """Initialize the shared attention and convolution backend."""
1032 super().__init__()
1033 # The target owns this executor, so keep its back-reference outside the module tree.
1034 object.__setattr__(self, "module", module)
1035 backend = backend.lower()
1036 if backend not in _GDN_BACKENDS:
1037 raise ValueError(
1038 f"unsupported GDN backend {backend!r}; "
1039 f"expected one of {sorted(_GDN_BACKENDS)}."
1040 )
1041 if chunk_size <= 0:
1042 raise ValueError(f"chunk_size must be positive, got {chunk_size}.")
1043 if backend in ("triton", "ascendc") and chunk_size != 64:
1044 raise NotImplementedError("Fused GDN backends require chunk_size=64.")
1045 self.gdn_backend = backend
1046 self.chunk_size = chunk_size
1047 self.cp_mesh = None if device_mesh is None else _ensure_1d(device_mesh)
1048 self.cp_size = 1 if self.cp_mesh is None else self.cp_mesh.size()
1049 self.cp_rank = 0 if self.cp_mesh is None else self.cp_mesh.get_local_rank()
1050 self.seq_dim = 1
1051 self.head_dim = 2
1052 self._validate_module()
1053 if self.gdn_backend != "eager":
1054 _validate_gdn_activation(self.module)
1055 self._validate_fused_conv()
1057 def _validate_module(self) -> None:
1058 """Validate the wrapped GDN module and Ulysses layout constraints."""
1059 required = (
1060 "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a",
1061 "conv1d", "A_log", "dt_bias", "norm", "out_proj",
1062 "num_k_heads", "num_v_heads", "head_k_dim", "head_v_dim",
1063 "key_dim", "value_dim", "conv_dim",
1064 )
1065 missing = [name for name in required if not hasattr(self.module, name)]
1066 if missing:
1067 raise TypeError(
1068 "Qwen3.5 GDN layer is missing required attributes: "
1069 + ", ".join(missing)
1070 )
1071 if self.module.num_k_heads % self.cp_size != 0:
1072 raise ValueError(
1073 f"linear attention num_k_heads ({self.module.num_k_heads}) must be "
1074 f"divisible by cp_size ({self.cp_size}) for Ulysses CP."
1075 )
1076 if self.module.num_v_heads % self.cp_size != 0:
1077 raise ValueError(
1078 f"linear attention num_v_heads ({self.module.num_v_heads}) must be "
1079 f"divisible by cp_size ({self.cp_size}) for Ulysses CP."
1080 )
1081 if self.module.num_v_heads % self.module.num_k_heads:
1082 raise ValueError("GDN num_v_heads must be divisible by num_k_heads.")
1083 _validate_gdn_activation(self.module)
1085 def _validate_fused_conv(self) -> None:
1086 """Reject source convolutions whose semantics differ from a causal depthwise kernel."""
1087 convolution = self.module.conv1d
1088 if convolution.stride != (1,) or convolution.dilation != (1,):
1089 raise ValueError("GDN fused convolution requires stride=1 and dilation=1.")
1090 if convolution.groups != convolution.in_channels or convolution.out_channels != convolution.in_channels:
1091 raise ValueError("GDN fused convolution requires equal input/output channels and depthwise groups.")
1092 if convolution.padding != (convolution.kernel_size[0] - 1,) or convolution.padding_mode != "zeros":
1093 raise ValueError("GDN fused convolution requires causal zero padding.")
1095 def _seq_to_head(self, tensor: torch.Tensor) -> torch.Tensor:
1096 if self.cp_size == 1:
1097 return tensor
1098 return _differentiable_all_to_all_shard(
1099 tensor,
1100 self.cp_mesh,
1101 split_dim=self.head_dim,
1102 concat_dim=self.seq_dim,
1103 )
1105 def _head_to_seq(self, tensor: torch.Tensor) -> torch.Tensor:
1106 if self.cp_size == 1:
1107 return tensor
1108 return _differentiable_all_to_all_shard(
1109 tensor,
1110 self.cp_mesh,
1111 split_dim=self.seq_dim,
1112 concat_dim=self.head_dim,
1113 )
1115 def _seq_to_head_qkvba(
1116 self,
1117 q_proj: torch.Tensor,
1118 k_proj: torch.Tensor,
1119 v_proj: torch.Tensor,
1120 b: torch.Tensor,
1121 a: torch.Tensor,
1122 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
1123 """Pack Q/K/V/B/A by CP rank and run a single seq-to-head all-to-all."""
1124 if self.cp_size == 1:
1125 return q_proj, k_proj, v_proj, b, a
1127 base = self.module
1128 local_key_dim = base.key_dim // self.cp_size
1129 local_value_dim = base.value_dim // self.cp_size
1130 local_num_v_heads = base.num_v_heads // self.cp_size
1132 q_chunks = torch.split(q_proj, local_key_dim, dim=-1)
1133 k_chunks = torch.split(k_proj, local_key_dim, dim=-1)
1134 v_chunks = torch.split(v_proj, local_value_dim, dim=-1)
1135 b_chunks = torch.split(b, local_num_v_heads, dim=-1)
1136 a_chunks = torch.split(a, local_num_v_heads, dim=-1)
1137 rank_major_chunks = [
1138 torch.cat(chunks, dim=-1)
1139 for chunks in zip(q_chunks, k_chunks, v_chunks, b_chunks, a_chunks)
1140 ]
1141 packed = torch.cat(rank_major_chunks, dim=-1).contiguous()
1142 packed = self._seq_to_head(packed)
1143 return torch.split(
1144 packed,
1145 [
1146 local_key_dim,
1147 local_key_dim,
1148 local_value_dim,
1149 local_num_v_heads,
1150 local_num_v_heads,
1151 ],
1152 dim=-1,
1153 )
1155 def _local_conv_weight(self) -> torch.Tensor:
1156 return _slice_qkv_local_cp(
1157 self.module.conv1d.weight,
1158 key_dim=self.module.key_dim,
1159 value_dim=self.module.value_dim,
1160 dim=0,
1161 cp_rank=self.cp_rank,
1162 cp_size=self.cp_size,
1163 )
1165 def _local_conv_bias(self) -> Optional[torch.Tensor]:
1166 bias = self.module.conv1d.bias
1167 if bias is None:
1168 return None
1169 return _slice_qkv_local_cp(
1170 bias,
1171 key_dim=self.module.key_dim,
1172 value_dim=self.module.value_dim,
1173 dim=0,
1174 cp_rank=self.cp_rank,
1175 cp_size=self.cp_size,
1176 )
1178 def _convolve_exchanged_qkv(self, q_proj, k_proj, v_proj, exchanged_info, full_seq_len, local_conv_dim):
1179 """Apply head-local ShortConv with the restored document boundaries."""
1180 base = self.module
1181 mixed_qkv = torch.cat((q_proj, k_proj, v_proj), dim=-1)
1182 weight, bias = self._local_conv_weight(), self._local_conv_bias()
1183 if exchanged_info is not None:
1184 return packed_causal_conv1d(
1185 mixed_qkv, base.conv1d, exchanged_info,
1186 weight=weight, bias=bias, backend=self.gdn_backend,
1187 )
1188 if self.gdn_backend != "eager":
1189 return fused_causal_conv1d(
1190 mixed_qkv, weight, bias, backend=self.gdn_backend,
1191 dilation=base.conv1d.dilation[0],
1192 )
1193 conv_out = F.conv1d(
1194 input=mixed_qkv.transpose(1, 2),
1195 weight=weight,
1196 bias=bias,
1197 stride=base.conv1d.stride,
1198 padding=base.conv1d.padding,
1199 dilation=base.conv1d.dilation,
1200 groups=local_conv_dim,
1201 )
1202 return F.silu(conv_out[:, :, :full_seq_len]).transpose(1, 2)
1204 def forward(
1205 self,
1206 hidden_states: torch.Tensor,
1207 cache_params: Optional[object] = None,
1208 attention_mask: Optional[torch.Tensor] = None,
1209 **kwargs,
1210 ) -> torch.Tensor:
1211 """Run Gated DeltaNet with pure Ulysses context parallel."""
1212 if cache_params is not None:
1213 raise NotImplementedError("GDN Ulysses CP currently supports training without cache only.")
1214 hidden_states = _local_tensor_at_cp_boundary(hidden_states)
1215 is_flat_input = hidden_states.ndim == 2
1216 if is_flat_input:
1217 hidden_states = hidden_states.unsqueeze(0)
1218 sequence_info = resolve_packed_sequence(
1219 kwargs, hidden_states.shape[1], cp_rank=self.cp_rank, cp_size=self.cp_size)
1220 if sequence_info is not None and hidden_states.shape[0] != 1:
1221 raise ValueError("Packed GDN requires batch dimension one.")
1223 base = self.module
1224 hidden_states = _apply_local_padding_mask(hidden_states, attention_mask)
1226 bsz, local_seq_len, _ = hidden_states.shape
1227 mixed_qkv = base.in_proj_qkv(hidden_states)
1228 z = base.in_proj_z(hidden_states).reshape(
1229 bsz,
1230 local_seq_len,
1231 base.num_v_heads,
1232 base.head_v_dim,
1233 )
1234 b = base.in_proj_b(hidden_states)
1235 a = base.in_proj_a(hidden_states)
1237 q_proj, k_proj, v_proj = torch.split(
1238 mixed_qkv,
1239 [base.key_dim, base.key_dim, base.value_dim],
1240 dim=-1,
1241 )
1242 q_proj, k_proj, v_proj, b, a = self._seq_to_head_qkvba(
1243 q_proj, k_proj, v_proj, b, a
1244 )
1246 full_seq_len = q_proj.shape[1]
1247 local_key_dim = base.key_dim // self.cp_size
1248 local_value_dim = base.value_dim // self.cp_size
1249 local_num_k_heads = base.num_k_heads // self.cp_size
1250 local_num_v_heads = base.num_v_heads // self.cp_size
1251 local_conv_dim = local_key_dim * 2 + local_value_dim
1253 exchanged_info = None if sequence_info is None else sequence_info.ulysses(self.cp_size, self.cp_rank)
1254 mixed_qkv = self._convolve_exchanged_qkv(
1255 q_proj, k_proj, v_proj, exchanged_info, full_seq_len, local_conv_dim,
1256 )
1258 query, key, value = torch.split(
1259 mixed_qkv,
1260 [local_key_dim, local_key_dim, local_value_dim],
1261 dim=-1,
1262 )
1263 query = query.reshape(bsz, full_seq_len, local_num_k_heads, base.head_k_dim)
1264 key = key.reshape(bsz, full_seq_len, local_num_k_heads, base.head_k_dim)
1265 value = value.reshape(bsz, full_seq_len, local_num_v_heads, base.head_v_dim)
1267 a_log = _slice_local_cp(base.A_log, 0, self.cp_rank, self.cp_size)
1268 dt_bias = _slice_local_cp(base.dt_bias, 0, self.cp_rank, self.cp_size)
1269 beta = b.sigmoid()
1270 g = -a_log.float().exp() * F.softplus(a.float() + dt_bias)
1272 kv_groups = _kv_groups(base)
1273 if kv_groups > 1:
1274 query = query.repeat_interleave(kv_groups, dim=2)
1275 key = key.repeat_interleave(kv_groups, dim=2)
1277 cu_seqlens = None
1278 if exchanged_info is not None:
1279 cu_seqlens = (exchanged_info.local_offsets if self.gdn_backend == "ascendc" else
1280 exchanged_info.materialize(query.device, self.chunk_size).cu_seqlens)
1281 core_attn_out, _ = chunk_gated_delta_rule(
1282 query,
1283 key,
1284 value,
1285 g=g,
1286 beta=beta,
1287 initial_state=None,
1288 output_final_state=False,
1289 chunk_size=self.chunk_size,
1290 use_qk_l2norm_in_kernel=True,
1291 backend=self.gdn_backend,
1292 cu_seqlens=cu_seqlens,
1293 )
1295 core_attn_out = self._head_to_seq(core_attn_out)
1296 core_attn_out = core_attn_out.reshape(-1, base.head_v_dim)
1297 z_flat = z.reshape(-1, base.head_v_dim)
1298 core_attn_out = base.norm(core_attn_out, z_flat)
1299 core_attn_out = core_attn_out.reshape(bsz, local_seq_len, base.value_dim)
1300 if hasattr(base, "out_proj_input"):
1301 core_attn_out = base.out_proj_input(core_attn_out)
1302 output = base.out_proj(core_attn_out)
1303 return output.squeeze(0) if is_flat_input else output
1306class GatedDeltaNetP2PCP(nn.Module):
1307 """Sequence-sharded GDN CP with an affine-summary state wavefront."""
1309 def __init__(
1310 self,
1311 module: nn.Module,
1312 device_mesh: DeviceMesh,
1313 *,
1314 backend: str = "eager",
1315 chunk_size: int = 64,
1316 ):
1317 super().__init__()
1318 # The target owns this executor, so keep its back-reference outside the module tree.
1319 object.__setattr__(self, "module", module)
1320 backend = backend.lower()
1321 if backend == "ascendc":
1322 raise NotImplementedError("GDN Ascend C state P2P is not supported; select the pure Ulysses wrapper.")
1323 if backend not in _GDN_BACKENDS:
1324 raise ValueError(
1325 f"unsupported GDN backend {backend!r}; "
1326 f"expected one of {sorted(_GDN_BACKENDS)}."
1327 )
1328 if chunk_size <= 0:
1329 raise ValueError(f"chunk_size must be positive, got {chunk_size}.")
1330 if backend == "triton" and chunk_size != 64:
1331 raise NotImplementedError("GDN Triton backend requires chunk_size=64.")
1332 self.gdn_backend = backend
1333 self.chunk_size = chunk_size
1334 self.cp_mesh = None if device_mesh is None else _ensure_1d(device_mesh)
1335 self.cp_size = 1 if self.cp_mesh is None else self.cp_mesh.size()
1336 self.cp_rank = 0 if self.cp_mesh is None else self.cp_mesh.get_local_rank()
1337 self._validate_module()
1339 def _validate_module(self) -> None:
1340 """Validate the Conv1d requirements of the P2P CP path."""
1341 required = (
1342 "in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a",
1343 "conv1d", "A_log", "dt_bias", "norm", "out_proj",
1344 "num_k_heads", "num_v_heads", "head_k_dim", "head_v_dim",
1345 "key_dim", "value_dim", "conv_dim",
1346 )
1347 missing = [name for name in required if not hasattr(self.module, name)]
1348 if missing:
1349 raise TypeError(
1350 "Qwen3.5 GDN layer is missing required attributes: "
1351 + ", ".join(missing)
1352 )
1353 if self.gdn_backend == "triton" and (
1354 self.module.head_k_dim != 128 or self.module.head_v_dim != 128
1355 ):
1356 raise NotImplementedError(
1357 "linear attention P2P Triton backend requires "
1358 "head_k_dim=head_v_dim=128."
1359 )
1360 if self.module.num_v_heads % self.module.num_k_heads:
1361 raise ValueError("GDN num_v_heads must be divisible by num_k_heads.")
1362 _validate_gdn_activation(self.module)
1363 conv = self.module.conv1d
1364 if conv.stride != (1,):
1365 raise ValueError(
1366 "linear attention P2P CP currently supports only conv1d stride=1."
1367 )
1368 if conv.groups != self.module.conv_dim:
1369 raise ValueError(
1370 "linear attention P2P CP expects depthwise conv1d groups=conv_dim."
1371 )
1372 if (
1373 conv.in_channels != self.module.conv_dim
1374 or conv.out_channels != self.module.conv_dim
1375 ):
1376 raise ValueError(
1377 "linear attention P2P CP expects conv1d channels to match conv_dim."
1378 )
1380 def _conv1d_with_halo(
1381 self, mixed_qkv: torch.Tensor, sequence_info: Optional[PackedSequenceInfo] = None,
1382 ) -> torch.Tensor:
1383 """Run local Conv1d after exchanging only the previous-rank halo."""
1384 return _causal_conv1d_with_cp_halo(
1385 mixed_qkv,
1386 self.module.conv1d,
1387 self.cp_mesh,
1388 self.cp_rank,
1389 self.cp_size,
1390 sequence_info=sequence_info,
1391 )
1393 def forward(
1394 self,
1395 hidden_states: torch.Tensor,
1396 cache_params: Optional[object] = None,
1397 attention_mask: Optional[torch.Tensor] = None,
1398 **kwargs,
1399 ) -> torch.Tensor:
1400 """Run Gated DeltaNet on local sequence shards with recurrent-state P2P."""
1401 if cache_params is not None:
1402 raise NotImplementedError("GDN P2P CP currently supports training without cache only.")
1403 hidden_states = _local_tensor_at_cp_boundary(hidden_states)
1404 is_flat_input = hidden_states.ndim == 2
1405 if is_flat_input:
1406 hidden_states = hidden_states.unsqueeze(0)
1407 sequence_info = resolve_packed_sequence(
1408 kwargs, hidden_states.shape[1], cp_rank=self.cp_rank, cp_size=self.cp_size)
1409 if sequence_info is not None and hidden_states.shape[0] != 1:
1410 raise ValueError("Packed GDN requires batch dimension one.")
1412 base = self.module
1413 hidden_states = _apply_local_padding_mask(hidden_states, attention_mask)
1415 bsz, local_seq_len, _ = hidden_states.shape
1416 mixed_qkv = base.in_proj_qkv(hidden_states)
1417 z = base.in_proj_z(hidden_states).reshape(
1418 bsz,
1419 local_seq_len,
1420 base.num_v_heads,
1421 base.head_v_dim,
1422 )
1423 b = base.in_proj_b(hidden_states)
1424 a = base.in_proj_a(hidden_states)
1426 mixed_qkv = self._conv1d_with_halo(mixed_qkv, sequence_info=sequence_info)
1427 query, key, value = torch.split(
1428 mixed_qkv,
1429 [base.key_dim, base.key_dim, base.value_dim],
1430 dim=-1,
1431 )
1432 query = query.reshape(bsz, local_seq_len, base.num_k_heads, base.head_k_dim)
1433 key = key.reshape(bsz, local_seq_len, base.num_k_heads, base.head_k_dim)
1434 value = value.reshape(bsz, local_seq_len, base.num_v_heads, base.head_v_dim)
1436 beta = b.sigmoid()
1437 g = -base.A_log.float().exp() * F.softplus(a.float() + base.dt_bias)
1439 kv_groups = _kv_groups(base)
1440 if kv_groups > 1:
1441 query = query.repeat_interleave(kv_groups, dim=2)
1442 key = key.repeat_interleave(kv_groups, dim=2)
1444 if self.gdn_backend == "triton":
1445 if not is_triton_gdn_available(query, key, value, g, beta):
1446 raise RuntimeError(
1447 "linear attention P2P Triton backend requires an NPU "
1448 "input satisfying the fixed GDN contract and a validated "
1449 "triton-ascend >= 3.2.1 installation."
1450 )
1451 core_attn_out = _gdn_state_p2p_triton(
1452 query,
1453 key,
1454 value,
1455 g,
1456 beta,
1457 self.cp_mesh,
1458 self.cp_rank,
1459 self.cp_size,
1460 sequence_info=sequence_info,
1461 )
1462 else:
1463 core_attn_out = _gdn_state_p2p_summary(
1464 query,
1465 key,
1466 value,
1467 g,
1468 beta,
1469 self.cp_mesh,
1470 self.cp_rank,
1471 self.cp_size,
1472 chunk_size=self.chunk_size,
1473 use_qk_l2norm_in_kernel=True,
1474 sequence_info=sequence_info,
1475 )
1477 core_attn_out = core_attn_out.reshape(-1, base.head_v_dim)
1478 z_flat = z.reshape(-1, base.head_v_dim)
1479 core_attn_out = base.norm(core_attn_out, z_flat)
1480 core_attn_out = core_attn_out.reshape(bsz, local_seq_len, base.value_dim)
1481 if hasattr(base, "out_proj_input"):
1482 core_attn_out = base.out_proj_input(core_attn_out)
1483 output = base.out_proj(core_attn_out)
1484 return output.squeeze(0) if is_flat_input else output
1487def _exchange_previous_cp_halo(tail: torch.Tensor, cp_mesh: Any) -> torch.Tensor:
1488 """Adapt channel-first convolution boundaries to the shared halo exchange."""
1489 halo = _all_to_all_previous_rank_halo(
1490 tail.transpose(1, 2).contiguous(), cp_mesh, cp_mesh.get_local_rank(), cp_mesh.size())
1491 return halo.transpose(1, 2).contiguous()
1494def _gdn_cp_causal_conv1d(
1495 original_conv: Callable[..., Any], cp_mesh: Any,
1496 hidden_states: torch.Tensor, weight: torch.Tensor,
1497 bias: Any = None, activation: Any = None, **kwargs: Any) -> Any:
1498 """Run the original causal Conv1d with the previous CP rank's halo."""
1499 halo_width = weight.shape[-1] - 1
1500 if halo_width <= 0 or cp_mesh.size() <= 1:
1501 return original_conv(
1502 hidden_states, weight, bias, activation, **kwargs)
1503 if hidden_states.shape[-1] < halo_width:
1504 raise ValueError(
1505 "GDN CP requires local sequence length "
1506 f"({hidden_states.shape[-1]}) to be at least Conv1d halo width "
1507 f"({halo_width})"
1508 )
1509 tail = hidden_states[..., -halo_width:].contiguous()
1510 halo = _exchange_previous_cp_halo(tail, cp_mesh)
1511 conv_input = torch.cat((halo, hidden_states), dim=-1)
1512 output = original_conv(
1513 conv_input, weight, bias, activation, **kwargs)
1514 return output[..., halo_width:].contiguous()
1517def _gdn_cp_conv1d_module(
1518 original_forward: Callable[..., Any], conv_module: Any,
1519 cp_mesh: Any, hidden_states: torch.Tensor) -> torch.Tensor:
1520 """Run the Transformers Conv1d fallback with the previous CP halo."""
1521 kernel_size = conv_module.kernel_size[0]
1522 dilation = conv_module.dilation[0]
1523 halo_width = (kernel_size - 1) * dilation
1524 if halo_width <= 0 or cp_mesh.size() <= 1:
1525 return original_forward(hidden_states)
1526 if hidden_states.shape[-1] < halo_width:
1527 raise ValueError(
1528 "GDN CP requires local sequence length "
1529 f"({hidden_states.shape[-1]}) to be at least Conv1d halo width "
1530 f"({halo_width})"
1531 )
1532 tail = hidden_states[..., -halo_width:].contiguous()
1533 halo = _exchange_previous_cp_halo(tail, cp_mesh)
1534 conv_input = torch.cat((halo, hidden_states), dim=-1)
1535 output = original_forward(conv_input)
1536 local_length = hidden_states.shape[-1]
1537 return output[..., halo_width:halo_width + local_length].contiguous()
1540def _gdn_rule_cp_to_hp(
1541 query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
1542 decay: torch.Tensor, beta: torch.Tensor,
1543 cp_mesh: Any) -> tuple[torch.Tensor, ...]:
1544 """Convert local-sequence GDN rule inputs into local-head layouts."""
1545 tensors = (query, key, value, decay, beta)
1546 expected_dims = (4, 4, 4, 3, 3)
1547 if any(
1548 tensor.dim() != expected
1549 for tensor, expected in zip(tensors, expected_dims)):
1550 shapes = [tuple(tensor.shape) for tensor in tensors]
1551 raise ValueError(
1552 "GDN CP expects 4-D Q/K/V and 3-D g/beta, got "
1553 f"{shapes}"
1554 )
1555 if any(tensor.shape[1] != query.shape[1] for tensor in tensors[1:]):
1556 raise ValueError("GDN CP inputs must share the sequence length")
1557 return tuple(
1558 ulysses_seq_to_head(
1559 tensor, seq_dim=1, head_dim=2, cp_mesh=cp_mesh)
1560 for tensor in tensors
1561 )
1564def _gdn_rule_hp_to_cp(
1565 output: torch.Tensor, cp_mesh: Any) -> torch.Tensor:
1566 """Restore one GDN rule output to the local-sequence layout."""
1567 if output.dim() != 4:
1568 raise ValueError(
1569 "GDN CP expects a 4-D GDN rule output, got "
1570 f"shape {tuple(output.shape)}"
1571 )
1572 return ulysses_head_to_seq(
1573 output, seq_dim=1, head_dim=2, cp_mesh=cp_mesh)
1576def _validate_gdn_cp(target_module: Any, cp_mesh: Any) -> None:
1577 """Validate the GDN interface required by CP interception."""
1578 if cp_mesh is None or cp_mesh.size() <= 1:
1579 raise ValueError("GDN CP requires an active CP mesh")
1580 module_name = type(target_module).__name__
1581 if "GatedDeltaNet" not in module_name:
1582 raise TypeError(
1583 "GDN CP wrapper requires a GatedDeltaNet module, got "
1584 f"{module_name}"
1585 )
1586 num_value_heads = getattr(target_module, "num_v_heads", None)
1587 if not isinstance(num_value_heads, int) or num_value_heads <= 0:
1588 raise ValueError(
1589 "GDN CP requires target_module.num_v_heads to be a "
1590 "positive integer"
1591 )
1592 if num_value_heads % cp_mesh.size():
1593 raise ValueError(
1594 f"GDN value heads ({num_value_heads}) must be divisible "
1595 f"by CP size ({cp_mesh.size()})"
1596 )
1599def _resolve_gdn_forward_implementation(
1600 original_forward: Callable[..., Any]) -> Callable[..., Any]:
1601 """Resolve the GDN implementation beneath Transformers hook wrappers."""
1602 supported_rule_names = {
1603 "torch_chunk_gated_delta_rule",
1604 "chunk_gated_delta_rule",
1605 }
1606 implementation = inspect.unwrap(original_forward)
1607 visited = set()
1608 while id(implementation) not in visited:
1609 visited.add(id(implementation))
1610 referenced_names = set(implementation.__code__.co_names)
1611 if supported_rule_names & referenced_names:
1612 return implementation
1613 closure = inspect.getclosurevars(implementation).nonlocals
1614 wrapped = closure.get("forward_func")
1615 if not callable(wrapped):
1616 break
1617 implementation = inspect.unwrap(wrapped)
1618 raise RuntimeError(
1619 "GDN CP could not resolve the Transformers forward "
1620 "implementation"
1621 )
1624def install_gdn_ulysses_cp(
1625 target_module: Any,
1626 cp_mesh: Any,
1627) -> None:
1628 """Install primitive-level Ulysses CP around a GDN forward."""
1629 _validate_gdn_cp(target_module, cp_mesh)
1630 original_forward = target_module.forward
1631 forward_impl = _resolve_gdn_forward_implementation(original_forward)
1632 forward_globals = forward_impl.__globals__
1633 referenced_names = set(forward_impl.__code__.co_names)
1635 conv_name = "causal_conv1d_fn"
1636 conv_owner = target_module if hasattr(target_module, conv_name) else forward_globals
1637 original_conv = (
1638 getattr(conv_owner, conv_name)
1639 if conv_owner is target_module else conv_owner.get(conv_name)
1640 )
1641 conv_module = getattr(target_module, "conv1d", None)
1642 original_conv_forward = getattr(conv_module, "forward", None)
1643 use_fused_conv = conv_name in referenced_names and callable(original_conv)
1644 use_module_conv = (
1645 not use_fused_conv
1646 and "conv1d" in referenced_names
1647 and callable(original_conv_forward)
1648 )
1649 if not use_fused_conv and not use_module_conv:
1650 raise RuntimeError(
1651 "GDN CP requires a fused or module causal Conv1d path"
1652 )
1654 rule_names = tuple(
1655 name for name in (
1656 "torch_chunk_gated_delta_rule",
1657 "chunk_gated_delta_rule",
1658 )
1659 if name in referenced_names and (
1660 callable(getattr(target_module, name, None))
1661 or callable(forward_globals.get(name))
1662 )
1663 )
1664 if not rule_names:
1665 raise RuntimeError(
1666 "GDN CP could not find a supported chunk Gated Delta "
1667 "Rule call in the Transformers forward"
1668 )
1669 rule_owners = {
1670 name: target_module if callable(getattr(target_module, name, None))
1671 else forward_globals
1672 for name in rule_names
1673 }
1674 original_rules = {
1675 name: getattr(owner, name) if owner is target_module else owner[name]
1676 for name, owner in rule_owners.items()
1677 }
1679 def replace_primitive(owner: Any, name: str, value: Callable[..., Any]) -> None:
1680 """Replace one module attribute or forward global primitive."""
1681 if owner is target_module:
1682 setattr(owner, name, value)
1683 else:
1684 owner[name] = value
1686 def make_cp_rule(original_rule: Callable[..., Any]) -> Callable[..., Any]:
1687 """Build one CP-aware wrapper for a referenced GDN rule primitive."""
1688 @functools.wraps(original_rule)
1689 def cp_rule(
1690 query: torch.Tensor,
1691 key: torch.Tensor,
1692 value: torch.Tensor,
1693 *args: Any,
1694 **kwargs: Any,
1695 ) -> Any:
1696 if args:
1697 raise TypeError(
1698 "GDN CP expects g and beta to be passed by keyword"
1699 )
1700 decay = kwargs.get("g")
1701 beta = kwargs.get("beta")
1702 if not isinstance(decay, torch.Tensor) or not isinstance(beta, torch.Tensor):
1703 raise TypeError(
1704 "GDN CP requires Tensor keyword arguments g and beta"
1705 )
1706 if kwargs.get("initial_state") is not None or kwargs.get(
1707 "output_final_state", False):
1708 raise NotImplementedError(
1709 "GDN CP does not support recurrent cache state"
1710 )
1711 query, key, value, decay, beta = _gdn_rule_cp_to_hp(
1712 query, key, value, decay, beta, cp_mesh)
1713 call_kwargs = dict(kwargs)
1714 call_kwargs.update(g=decay, beta=beta)
1715 output = original_rule(query, key, value, **call_kwargs)
1716 if not isinstance(output, tuple) or len(output) != 2:
1717 raise TypeError(
1718 "GDN rule must return (output, recurrent_state)"
1719 )
1720 core_output, recurrent_state = output
1721 if recurrent_state is not None:
1722 raise NotImplementedError(
1723 "GDN CP does not support a returned recurrent state"
1724 )
1725 return _gdn_rule_hp_to_cp(core_output, cp_mesh), recurrent_state
1727 return cp_rule
1729 @functools.wraps(original_forward)
1730 def cp_forward(*args: Any, **kwargs: Any) -> Any:
1731 """Run the original GDN forward with CP-aware primitive interfaces."""
1732 cache_params = kwargs.get("cache_params")
1733 if cache_params is None and len(args) > 1:
1734 cache_params = args[1]
1735 if cache_params is not None:
1736 raise NotImplementedError(
1737 "GDN CP currently supports training with cache_params=None"
1738 )
1739 if kwargs.get("cu_seq_lens_q") is not None:
1740 raise NotImplementedError(
1741 "GDN CP does not support packed sequences"
1742 )
1744 fired = {"conv": 0, "rule": 0}
1746 if use_fused_conv:
1747 @functools.wraps(original_conv)
1748 def cp_conv(*conv_args: Any, **conv_kwargs: Any) -> Any:
1749 fired["conv"] += 1
1750 if not conv_args and "x" in conv_kwargs:
1751 conv_kwargs = dict(conv_kwargs)
1752 conv_args = (conv_kwargs.pop("x"),)
1753 return _gdn_cp_causal_conv1d(
1754 original_conv, cp_mesh, *conv_args, **conv_kwargs)
1755 else:
1756 @functools.wraps(original_conv_forward)
1757 def cp_conv(hidden_states: torch.Tensor) -> torch.Tensor:
1758 fired["conv"] += 1
1759 return _gdn_cp_conv1d_module(
1760 original_conv_forward, conv_module, cp_mesh, hidden_states)
1762 def count_rule(wrapped_rule):
1763 @functools.wraps(wrapped_rule)
1764 def counted_rule(*rule_args: Any, **rule_kwargs: Any) -> Any:
1765 fired["rule"] += 1
1766 return wrapped_rule(*rule_args, **rule_kwargs)
1768 return counted_rule
1770 wrapped_rules = {
1771 name: count_rule(make_cp_rule(original_rule))
1772 for name, original_rule in original_rules.items()
1773 }
1775 if use_fused_conv:
1776 replace_primitive(conv_owner, conv_name, cp_conv)
1777 else:
1778 conv_module.forward = cp_conv
1779 for name, wrapped_rule in wrapped_rules.items():
1780 replace_primitive(rule_owners[name], name, wrapped_rule)
1781 try:
1782 output = original_forward(*args, **kwargs)
1783 finally:
1784 if use_fused_conv:
1785 replace_primitive(conv_owner, conv_name, original_conv)
1786 else:
1787 conv_module.forward = original_conv_forward
1788 for name, original_rule in original_rules.items():
1789 replace_primitive(rule_owners[name], name, original_rule)
1790 if fired["conv"] != 1 or fired["rule"] != 1:
1791 raise RuntimeError(
1792 "GDN CP wrapper expected one causal Conv1d and one "
1793 "Gated Delta Rule call, got "
1794 f"conv={fired['conv']}, rule={fired['rule']}"
1795 )
1796 return output
1798 target_module.forward = cp_forward