wehub-resource-sync
614 lines
23 KiB
Python
614 lines
23 KiB
Python
"""Distributed GDN scan with Context Parallel state correction.
|
|
|
|
Each GPU runs a local scan on its T/P frames, then corrects the results
|
|
using the true initial state obtained via all-gather + merge.
|
|
|
|
Communication is O(P * D^2) via all-gather -- symmetric collective that
|
|
avoids cross-communicator deadlocks when FSDP and Ulysses SP operate on
|
|
other NCCL process groups concurrently.
|
|
|
|
Algorithm:
|
|
|
|
1. Local scan with S_init=0 --> S_local[t]
|
|
2. Cumulative transition products --> W_cum[t]
|
|
3. Extract chunk composites: h_ext = S_local[-1], M = W_cum[-1]
|
|
4. All-gather (h_ext, M) across P ranks
|
|
5. Merge: compose predecessors to get S_init
|
|
6. Correction: S_corrected[t] = f(S_init, W_cum[t]) + S_local[t]
|
|
|
|
The KV state uses right-multiply: S = S_prev @ W + U
|
|
The Z state uses left-multiply: S = W @ S_prev + U
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, NamedTuple
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
from torch import Tensor
|
|
from torch.distributed import ProcessGroup
|
|
|
|
|
|
class CpFrameGdnScanResult(NamedTuple):
|
|
"""Return type of :func:`cp_frame_gdn_scan` when ``truncate_to_active`` is set.
|
|
|
|
Carries the per-position corrected scan outputs (same shape as the
|
|
legacy 2-tuple return) plus the terminal recurrence state at logical
|
|
global position ``truncate_to_active - 1`` (identical on every rank,
|
|
broadcast from the owning rank).
|
|
|
|
Note: ``NamedTuple`` iterates ALL fields when unpacked. Callers that
|
|
use the legacy 2-tuple unpacking (``S_kv, S_z = cp_frame_gdn_scan(...)``)
|
|
MUST NOT pass ``truncate_to_active``; instead use the default
|
|
``truncate_to_active=None`` path which returns a plain 2-tuple.
|
|
"""
|
|
|
|
S_kv_all: Tensor # (BH, T_local, D, D)
|
|
S_z_all: Tensor # (BH, T_local, D)
|
|
terminal_state_kv: Tensor # (BH, D, D), same on every rank
|
|
terminal_state_z: Tensor # (BH, D), same on every rank
|
|
|
|
|
|
from diffusion.distributed.context_parallel.config import (
|
|
get_cp_allgather_impl,
|
|
get_cp_scan_backend,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Local scan backends
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@torch.compile(dynamic=True)
|
|
def _pytorch_scan_compiled(
|
|
W_kv: Tensor,
|
|
U_kv: Tensor,
|
|
W_z: Tensor,
|
|
U_z: Tensor,
|
|
S_init_kv: Tensor | None = None,
|
|
S_init_z: Tensor | None = None,
|
|
) -> tuple[Tensor, Tensor]:
|
|
"""Compiled local scan for CP.
|
|
|
|
``torch.compile`` traces through the loop and generates an efficient
|
|
fused kernel with automatic backward differentiation. All computation
|
|
is done in FP32 for numerical stability.
|
|
"""
|
|
orig_dtype = W_kv.dtype
|
|
W_kv, U_kv = W_kv.float(), U_kv.float()
|
|
W_z, U_z = W_z.float(), U_z.float()
|
|
|
|
BH, T, D, _ = W_kv.shape
|
|
if S_init_kv is not None:
|
|
S_kv = S_init_kv.float()
|
|
else:
|
|
S_kv = torch.zeros(BH, D, D, device=W_kv.device, dtype=torch.float32)
|
|
|
|
if S_init_z is not None:
|
|
S_z = S_init_z.float()
|
|
else:
|
|
S_z = torch.zeros(BH, D, device=U_z.device, dtype=torch.float32)
|
|
|
|
S_kv_all = torch.empty_like(U_kv)
|
|
S_z_all = torch.empty_like(U_z)
|
|
for t in range(T):
|
|
S_kv = torch.matmul(S_kv, W_kv[:, t]) + U_kv[:, t]
|
|
S_z = torch.bmm(W_z[:, t], S_z.unsqueeze(-1)).squeeze(-1) + U_z[:, t]
|
|
S_kv_all[:, t] = S_kv
|
|
S_z_all[:, t] = S_z
|
|
return S_kv_all.to(orig_dtype), S_z_all.to(orig_dtype)
|
|
|
|
|
|
class _PyTorchScan:
|
|
"""Wrapper that mimics the ``autograd.Function`` ``.apply()`` interface
|
|
while delegating to the compiled scan function.
|
|
|
|
``torch.compile`` handles backward differentiation automatically, so
|
|
a custom ``autograd.Function`` is no longer needed.
|
|
"""
|
|
|
|
@staticmethod
|
|
def apply(
|
|
W_kv: Tensor,
|
|
U_kv: Tensor,
|
|
W_z: Tensor,
|
|
U_z: Tensor,
|
|
S_init_kv: Tensor | None = None,
|
|
S_init_z: Tensor | None = None,
|
|
) -> tuple[Tensor, Tensor]:
|
|
return _pytorch_scan_compiled(W_kv, U_kv, W_z, U_z, S_init_kv, S_init_z)
|
|
|
|
|
|
def _get_local_scan_cls(device_is_cuda: bool) -> type:
|
|
"""Select the local scan implementation based on config.
|
|
|
|
Args:
|
|
device_is_cuda: Whether the tensors reside on a CUDA device.
|
|
|
|
Returns:
|
|
``_PyTorchScan`` or ``FrameGDNScan`` (Triton) autograd Function class.
|
|
"""
|
|
backend = get_cp_scan_backend()
|
|
if backend == "triton":
|
|
from diffusion.model.ops.frame_gdn.scan_triton import FrameGDNScan
|
|
|
|
return FrameGDNScan
|
|
return _PyTorchScan
|
|
|
|
|
|
# Keep backward-compatible alias.
|
|
get_local_scan_cls = _get_local_scan_cls
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Cumulative matrix products
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@torch.compile(dynamic=True)
|
|
def _cumulative_matmul_right(W: Tensor) -> Tensor:
|
|
"""Cumulative right-multiply: W_cum[t] = W[0] @ W[1] @ ... @ W[t].
|
|
|
|
Args:
|
|
W: ``(BH, T, D, D)``
|
|
|
|
Returns:
|
|
W_cum: ``(BH, T, D, D)`` where ``W_cum[:, t]`` is the cumulative
|
|
product of transition matrices up to and including step *t*.
|
|
"""
|
|
slices: list[Tensor] = [W[:, 0]]
|
|
for t in range(1, W.shape[1]):
|
|
slices.append(torch.matmul(slices[-1], W[:, t]))
|
|
return torch.stack(slices, dim=1)
|
|
|
|
|
|
@torch.compile(dynamic=True)
|
|
def _cumulative_matmul_left(W: Tensor) -> Tensor:
|
|
"""Cumulative left-multiply: W_cum[t] = W[t] @ ... @ W[1] @ W[0].
|
|
|
|
For the Z state with left-multiply convention.
|
|
|
|
Args:
|
|
W: ``(BH, T, D, D)``
|
|
|
|
Returns:
|
|
W_cum: ``(BH, T, D, D)``
|
|
"""
|
|
slices: list[Tensor] = [W[:, 0]]
|
|
for t in range(1, W.shape[1]):
|
|
slices.append(torch.matmul(W[:, t], slices[-1]))
|
|
return torch.stack(slices, dim=1)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# All-gather helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
from diffusion.distributed.context_parallel.halo_exchange import _to_global_rank
|
|
|
|
|
|
def _allgather(tensor: Tensor, group: ProcessGroup) -> Tensor:
|
|
"""All-gather a tensor across the group, returning ``(P, *shape)``."""
|
|
world = dist.get_world_size(group)
|
|
rank = dist.get_rank(group)
|
|
tensor_contig = tensor.contiguous()
|
|
out = torch.empty((world,) + tensor_contig.shape, dtype=tensor_contig.dtype, device=tensor_contig.device)
|
|
impl = get_cp_allgather_impl()
|
|
|
|
if impl == "collective":
|
|
# ``all_gather_into_tensor`` concatenates rank inputs along dim 0.
|
|
# Reshape back to the stacked ``(world, *shape)`` contract used by
|
|
# this module. This works for both Gloo and NCCL.
|
|
flat_out = torch.empty(
|
|
(world * tensor_contig.shape[0],) + tuple(tensor_contig.shape[1:]),
|
|
dtype=tensor_contig.dtype,
|
|
device=tensor_contig.device,
|
|
)
|
|
dist.all_gather_into_tensor(flat_out, tensor_contig, group=group)
|
|
out = flat_out.reshape((world,) + tuple(tensor_contig.shape))
|
|
elif impl == "list":
|
|
# Conservative fallback for communicator behavior checks.
|
|
gathered = [torch.empty_like(tensor_contig) for _ in range(world)]
|
|
dist.all_gather(gathered, tensor_contig, group=group)
|
|
out = torch.stack(gathered, dim=0)
|
|
else:
|
|
# FSDP2-oriented P2P implementation.
|
|
ops = []
|
|
out[rank].copy_(tensor_contig)
|
|
for i in range(world):
|
|
if i != rank:
|
|
peer = _to_global_rank(group, i)
|
|
ops.append(dist.P2POp(dist.isend, tensor_contig, peer, group=group))
|
|
ops.append(dist.P2POp(dist.irecv, out[i], peer, group=group))
|
|
|
|
if ops:
|
|
reqs = dist.batch_isend_irecv(ops)
|
|
for req in reqs:
|
|
req.wait()
|
|
|
|
return out
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# All-gather + merge autograd Function
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _CPAllGatherMerge(torch.autograd.Function):
|
|
"""Differentiable all-gather + exclusive prefix merge for CP.
|
|
|
|
Each rank contributes its chunk composite ``(h_ext, M)`` for both the
|
|
KV and Z scans. All composites are all-gathered, then each rank
|
|
locally computes the exclusive prefix composition of all preceding
|
|
chunks to obtain the correct initial state ``S_init``.
|
|
|
|
Handles two multiply conventions simultaneously:
|
|
- KV (right-multiply): S_final = S_init @ M + h_ext
|
|
- Z (left-multiply): S_final = M @ S_init + h_ext
|
|
|
|
Args (forward):
|
|
h_ext_kv: ``(BH, D, D)`` -- KV input composite (local scan final state).
|
|
M_kv: ``(BH, D, D)`` -- KV transition composite (cumulative product).
|
|
h_ext_z: ``(BH, D)`` -- Z input composite.
|
|
M_z: ``(BH, D, D)`` -- Z transition composite.
|
|
group: CP process group.
|
|
reverse: If True, state flows from rank P-1 to rank 0.
|
|
|
|
Returns:
|
|
S_init_kv: ``(BH, D, D)`` -- correct KV initial state for this rank.
|
|
S_init_z: ``(BH, D)`` -- correct Z initial state for this rank.
|
|
"""
|
|
|
|
@staticmethod
|
|
def forward(
|
|
ctx: Any,
|
|
h_ext_kv: Tensor,
|
|
M_kv: Tensor,
|
|
h_ext_z: Tensor,
|
|
M_z: Tensor,
|
|
group: ProcessGroup,
|
|
reverse: bool = False,
|
|
) -> tuple[Tensor, Tensor]:
|
|
rank = dist.get_rank(group)
|
|
world = dist.get_world_size(group)
|
|
|
|
# All-gather composites from all ranks.
|
|
h_all_kv = _allgather(h_ext_kv, group) # (P, BH, D, D)
|
|
M_all_kv = _allgather(M_kv, group) # (P, BH, D, D)
|
|
h_all_z = _allgather(h_ext_z, group) # (P, BH, D)
|
|
M_all_z = _allgather(M_z, group) # (P, BH, D, D)
|
|
|
|
if reverse:
|
|
logical_rank = world - 1 - rank
|
|
h_all_kv = h_all_kv.flip(0)
|
|
M_all_kv = M_all_kv.flip(0)
|
|
h_all_z = h_all_z.flip(0)
|
|
M_all_z = M_all_z.flip(0)
|
|
else:
|
|
logical_rank = rank
|
|
|
|
# Exclusive prefix composition: compose chunks 0..logical_rank-1.
|
|
S_init_kv, S_init_z = _exclusive_prefix_compose(
|
|
h_all_kv,
|
|
M_all_kv,
|
|
h_all_z,
|
|
M_all_z,
|
|
logical_rank,
|
|
)
|
|
|
|
ctx.save_for_backward(M_kv, M_z, S_init_kv, S_init_z)
|
|
ctx.group = group
|
|
ctx.reverse = reverse
|
|
ctx.world = world
|
|
ctx.rank = rank
|
|
return S_init_kv, S_init_z
|
|
|
|
@staticmethod
|
|
def backward(
|
|
ctx: Any,
|
|
dS_init_kv: Tensor,
|
|
dS_init_z: Tensor,
|
|
) -> tuple[Tensor, Tensor, Tensor, Tensor, None, None]:
|
|
M_kv, M_z, S_init_kv, S_init_z = ctx.saved_tensors
|
|
group = ctx.group
|
|
reverse = ctx.reverse
|
|
world = ctx.world
|
|
rank = ctx.rank
|
|
|
|
compute_dtype = torch.float32
|
|
dS_init_kv = dS_init_kv.to(compute_dtype)
|
|
dS_init_z = dS_init_z.to(compute_dtype)
|
|
M_kv = M_kv.to(compute_dtype)
|
|
M_z = M_z.to(compute_dtype)
|
|
S_init_kv = S_init_kv.to(compute_dtype)
|
|
S_init_z = S_init_z.to(compute_dtype)
|
|
|
|
if reverse:
|
|
logical_rank = world - 1 - rank
|
|
else:
|
|
logical_rank = rank
|
|
|
|
# All-gather dS_init and M from all ranks.
|
|
dS_all_kv = _allgather(dS_init_kv, group) # (P, BH, D, D)
|
|
dS_all_z = _allgather(dS_init_z, group) # (P, BH, D)
|
|
M_all_kv = _allgather(M_kv, group) # (P, BH, D, D)
|
|
M_all_z = _allgather(M_z, group) # (P, BH, D, D)
|
|
|
|
if reverse:
|
|
dS_all_kv = dS_all_kv.flip(0)
|
|
dS_all_z = dS_all_z.flip(0)
|
|
M_all_kv = M_all_kv.flip(0)
|
|
M_all_z = M_all_z.flip(0)
|
|
|
|
# Compute dS_final: gradient flowing into this rank's composite
|
|
# from all successor ranks.
|
|
#
|
|
# In the forward, rank j > logical_rank uses our (h, M) through:
|
|
# S_init(j) = ... @ M[logical_rank] + h[logical_rank] ...
|
|
#
|
|
# The backward sweep accumulates:
|
|
# sent[r] = dS_init[r] + sent[r+1] @ M[r]^T (KV, right-multiply)
|
|
# sent[r] = dS_init[r] + M[r]^T @ sent[r+1] (Z, left-multiply)
|
|
# and dS_final[logical_rank] = sent[logical_rank + 1].
|
|
if logical_rank >= world - 1:
|
|
dS_final_kv = torch.zeros_like(dS_init_kv)
|
|
dS_final_z = torch.zeros_like(dS_init_z)
|
|
else:
|
|
sent_kv = dS_all_kv[world - 1].clone()
|
|
sent_z = dS_all_z[world - 1].clone()
|
|
for r in range(world - 2, logical_rank, -1):
|
|
sent_kv = dS_all_kv[r] + torch.matmul(
|
|
sent_kv,
|
|
M_all_kv[r].transpose(-1, -2),
|
|
)
|
|
sent_z = dS_all_z[r] + torch.bmm(
|
|
M_all_z[r].transpose(-1, -2),
|
|
sent_z.unsqueeze(-1),
|
|
).squeeze(-1)
|
|
dS_final_kv = sent_kv
|
|
dS_final_z = sent_z
|
|
|
|
# Gradients w.r.t. this rank's composites.
|
|
# Forward: S_final = S_init @ M + h_ext (KV)
|
|
# S_final = M @ S_init + h_ext (Z)
|
|
dh_ext_kv = dS_final_kv
|
|
dM_kv = torch.matmul(S_init_kv.transpose(-1, -2), dS_final_kv)
|
|
|
|
dh_ext_z = dS_final_z
|
|
dM_z = torch.bmm(dS_final_z.unsqueeze(-1), S_init_z.unsqueeze(-2))
|
|
|
|
return dh_ext_kv, dM_kv, dh_ext_z, dM_z, None, None
|
|
|
|
|
|
class _BroadcastFromLastRank(torch.autograd.Function):
|
|
"""Autograd-aware broadcast of a tensor from the LAST CP rank.
|
|
|
|
Forward:
|
|
Every rank emits the value of ``tensor`` from the last rank (the
|
|
non-last ranks' input value is DROPPED). Equivalent to
|
|
``dist.broadcast(tensor, src=last_rank)`` but autograd-tracked.
|
|
|
|
Backward:
|
|
Gradient flowing into the broadcasted output on EVERY rank is
|
|
summed (all-reduced) and accumulated into the last rank's source
|
|
tensor. Non-last ranks receive zero gradient (they didn't
|
|
contribute to the forward).
|
|
|
|
This is mathematically equivalent to a "scatter" of the source value
|
|
to every rank with a "sum" gradient back to the source.
|
|
"""
|
|
|
|
@staticmethod
|
|
def forward(ctx: Any, tensor: Tensor, group: ProcessGroup) -> Tensor:
|
|
rank = dist.get_rank(group)
|
|
world = dist.get_world_size(group)
|
|
last_rank_global = _to_global_rank(group, world - 1)
|
|
ctx.group = group
|
|
ctx.world = world
|
|
ctx.rank = rank
|
|
ctx.last_rank_global = last_rank_global
|
|
|
|
out = tensor.detach().clone().contiguous()
|
|
# Single broadcast: out becomes the last rank's value on every rank.
|
|
if world > 1:
|
|
dist.broadcast(out, src=last_rank_global, group=group)
|
|
return out
|
|
|
|
@staticmethod
|
|
def backward(ctx: Any, grad_out: Tensor) -> tuple[Tensor, None]:
|
|
group = ctx.group
|
|
world = ctx.world
|
|
rank = ctx.rank
|
|
|
|
if world <= 1:
|
|
return grad_out, None
|
|
|
|
# Sum gradients across ranks. The result is the total gradient flowing
|
|
# into the source (last rank's) terminal state. Only the last rank
|
|
# returns this sum; non-last ranks return zeros (their input was
|
|
# dropped in the forward).
|
|
summed = grad_out.contiguous().clone()
|
|
dist.all_reduce(summed, op=dist.ReduceOp.SUM, group=group)
|
|
if rank == world - 1:
|
|
return summed, None
|
|
return torch.zeros_like(grad_out), None
|
|
|
|
|
|
def _exclusive_prefix_compose(
|
|
h_all_kv: Tensor,
|
|
M_all_kv: Tensor,
|
|
h_all_z: Tensor,
|
|
M_all_z: Tensor,
|
|
logical_rank: int,
|
|
) -> tuple[Tensor, Tensor]:
|
|
"""Compose chunks 0, 1, ..., logical_rank-1 to get S_init.
|
|
|
|
For logical_rank == 0, returns zeros (first rank starts from zero state).
|
|
|
|
KV (right-multiply): h = h @ M[j] + h_ext[j] for j = 0..rank-1
|
|
Z (left-multiply): h = M[j] @ h + h_ext[j] for j = 0..rank-1
|
|
"""
|
|
if logical_rank == 0:
|
|
return torch.zeros_like(h_all_kv[0]), torch.zeros_like(h_all_z[0])
|
|
|
|
S_kv = torch.zeros_like(h_all_kv[0])
|
|
S_z = torch.zeros_like(h_all_z[0])
|
|
for j in range(logical_rank):
|
|
S_kv = torch.matmul(S_kv, M_all_kv[j]) + h_all_kv[j]
|
|
S_z = torch.bmm(M_all_z[j], S_z.unsqueeze(-1)).squeeze(-1) + h_all_z[j]
|
|
return S_kv, S_z
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Public API
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def cp_frame_gdn_scan(
|
|
W_kv: Tensor,
|
|
U_kv: Tensor,
|
|
W_z: Tensor,
|
|
U_z: Tensor,
|
|
group: ProcessGroup,
|
|
reverse: bool = False,
|
|
truncate_to_active: int | None = None,
|
|
) -> tuple[Tensor, Tensor] | CpFrameGdnScanResult:
|
|
"""Distributed GDN scan across CP ranks with state correction.
|
|
|
|
Produces the same result as running the scan on the globally concatenated
|
|
(W, U) sequence, but each GPU only touches its local T_local frames plus
|
|
O(P * D^2) communication via all-gather.
|
|
|
|
Algorithm:
|
|
1. Local scan with S_init = 0
|
|
2. Cumulative transition products
|
|
3. Extract chunk composites (h_ext, M)
|
|
4. All-gather + merge to get S_init
|
|
5. Correct all local states
|
|
|
|
Args:
|
|
W_kv: ``(BH, T_local, D, D)`` -- local KV transition matrices.
|
|
U_kv: ``(BH, T_local, D, D)`` -- local KV input matrices.
|
|
W_z: ``(BH, T_local, D, D)`` -- local Z transition matrices.
|
|
U_z: ``(BH, T_local, D)`` -- local Z input vectors.
|
|
group: CP process group.
|
|
reverse: If True, the scan direction is reversed (for backward
|
|
recurrence in BidirectionalGDN).
|
|
truncate_to_active: When set to an integer
|
|
``K_active`` (logical valid global cond length), the scan
|
|
internally masks ``(W, U)`` at positions ``>= K_active`` so
|
|
those positions do NOT contribute to state propagation
|
|
(``W = I``, ``U = 0``). The scan also extracts the terminal
|
|
state at global position ``K_active - 1`` and broadcasts it
|
|
to every CP rank. Return shape changes to
|
|
:class:`CpFrameGdnScanResult` (NamedTuple of 4 fields).
|
|
|
|
Constraints (forward direction only, ``reverse=False``):
|
|
``1 <= K_active <= T_local * cp_size``.
|
|
``reverse=True`` is not supported (no AR rollout consumer).
|
|
``cp_size=1`` is supported (the mask still applies; terminal
|
|
state is extracted locally without communication).
|
|
|
|
Returns:
|
|
When ``truncate_to_active is None`` (default, backward-compatible
|
|
path): plain 2-tuple
|
|
``(S_kv_all: (BH, T_local, D, D), S_z_all: (BH, T_local, D))``.
|
|
|
|
When ``truncate_to_active`` is set: :class:`CpFrameGdnScanResult`
|
|
with fields ``S_kv_all``, ``S_z_all``, ``terminal_state_kv``
|
|
``(BH, D, D)``, ``terminal_state_z`` ``(BH, D)``. Terminal state
|
|
is identical on every CP rank.
|
|
"""
|
|
# Handle truncate_to_active by masking W/U at padded positions so state
|
|
# propagation stops at position ``K_active - 1``.
|
|
if truncate_to_active is not None:
|
|
if reverse:
|
|
raise NotImplementedError(
|
|
"cp_frame_gdn_scan: truncate_to_active is only supported for "
|
|
"reverse=False (no AR rollout consumer needs reverse trunc)."
|
|
)
|
|
T_local_in = W_kv.shape[1]
|
|
D_dim = W_kv.shape[-1]
|
|
cp_world = dist.get_world_size(group)
|
|
cp_rank_in = dist.get_rank(group)
|
|
T_global = T_local_in * cp_world
|
|
K_active = int(truncate_to_active)
|
|
if K_active < 1 or K_active > T_global:
|
|
raise ValueError(f"truncate_to_active={K_active} must satisfy 1 <= K_active " f"<= T_global={T_global}")
|
|
|
|
# Build per-rank position mask: positions >= K_active should have
|
|
# W=I and U=0. The mask is ``valid[local_t] = (rank * T_local + local_t < K_active)``.
|
|
local_positions = torch.arange(T_local_in, device=W_kv.device)
|
|
global_positions = cp_rank_in * T_local_in + local_positions
|
|
valid_mask = (global_positions < K_active).to(W_kv.dtype) # (T_local,) 0/1
|
|
valid_kv = valid_mask.view(1, T_local_in, 1, 1) # (1, T_local, 1, 1)
|
|
valid_z_W = valid_mask.view(1, T_local_in, 1, 1)
|
|
valid_z_U = valid_mask.view(1, T_local_in, 1)
|
|
|
|
eye_kv = torch.eye(D_dim, device=W_kv.device, dtype=W_kv.dtype).view(1, 1, D_dim, D_dim)
|
|
eye_z = torch.eye(D_dim, device=W_z.device, dtype=W_z.dtype).view(1, 1, D_dim, D_dim)
|
|
|
|
# W -> I, U -> 0 at padded positions.
|
|
W_kv = valid_kv * W_kv + (1.0 - valid_kv) * eye_kv
|
|
U_kv = valid_kv * U_kv # 0 at padded.
|
|
W_z = valid_z_W * W_z + (1.0 - valid_z_W) * eye_z
|
|
U_z = valid_z_U * U_z
|
|
|
|
# --- Step 1: cumulative transition products ---
|
|
W_kv_cum = _cumulative_matmul_right(W_kv) # (BH, T_local, D, D)
|
|
W_z_cum = _cumulative_matmul_left(W_z) # (BH, T_local, D, D)
|
|
|
|
# --- Step 2: local scan with S_init = 0 to get h_ext ---
|
|
local_scan = _get_local_scan_cls(W_kv.is_cuda)
|
|
S_kv_local, S_z_local = local_scan.apply(W_kv, U_kv, W_z, U_z)
|
|
|
|
# --- Step 3: extract chunk composites ---
|
|
h_ext_kv = S_kv_local[:, -1] # (BH, D, D)
|
|
M_kv = W_kv_cum[:, -1] # (BH, D, D)
|
|
h_ext_z = S_z_local[:, -1] # (BH, D)
|
|
M_z = W_z_cum[:, -1] # (BH, D, D)
|
|
|
|
# --- Step 4: all-gather + merge to get correct S_init ---
|
|
S_init_kv, S_init_z = _CPAllGatherMerge.apply(
|
|
h_ext_kv,
|
|
M_kv,
|
|
h_ext_z,
|
|
M_z,
|
|
group,
|
|
reverse,
|
|
)
|
|
|
|
# --- Step 5: additive correction (replaces full rescan) ---
|
|
# By linearity of the recurrence s[t] = s[t-1] @ W[t] + U[t]:
|
|
# S_corrected[t] = S_zero[t] + S_init @ W_cum[t]
|
|
# This is a parallel matmul instead of a sequential scan.
|
|
# KV (right-multiply convention): S[t] = S[t-1] @ W[t] + U[t]
|
|
S_kv_corrected = S_kv_local + torch.matmul(S_init_kv.unsqueeze(1), W_kv_cum)
|
|
# Z (left-multiply convention): S[t] = W[t] @ S[t-1] + U[t]
|
|
# W_z_cum[t] = W[t] @ ... @ W[0], so correction = W_z_cum[t] @ S_init_z
|
|
T_local = W_z_cum.shape[1]
|
|
S_z_corrected = S_z_local + torch.bmm(
|
|
W_z_cum.reshape(-1, W_z_cum.shape[2], W_z_cum.shape[3]),
|
|
S_init_z.unsqueeze(1).expand(-1, T_local, -1).reshape(-1, W_z_cum.shape[3], 1),
|
|
).reshape(S_z_local.shape)
|
|
|
|
if truncate_to_active is None:
|
|
return S_kv_corrected, S_z_corrected
|
|
|
|
# Extract terminal state at global position ``K_active - 1`` and
|
|
# broadcast to all ranks. Because padded positions have W=I, U=0, the
|
|
# recurrence state stays constant after the active prefix.
|
|
terminal_kv_local = S_kv_corrected[:, -1].contiguous() # (BH, D, D)
|
|
terminal_z_local = S_z_corrected[:, -1].contiguous() # (BH, D)
|
|
terminal_kv = _BroadcastFromLastRank.apply(terminal_kv_local, group)
|
|
terminal_z = _BroadcastFromLastRank.apply(terminal_z_local, group)
|
|
|
|
return CpFrameGdnScanResult(
|
|
S_kv_all=S_kv_corrected,
|
|
S_z_all=S_z_corrected,
|
|
terminal_state_kv=terminal_kv,
|
|
terminal_state_z=terminal_z,
|
|
)
|